Skip to content

Meta #426: MMM generative workflows as stress-test for panel + simulate APIs - #468

Open
drbenvincent wants to merge 10 commits into
mainfrom
run/426-mmm-simulate-stress-test
Open

Meta #426: MMM generative workflows as stress-test for panel + simulate APIs#468
drbenvincent wants to merge 10 commits into
mainfrom
run/426-mmm-simulate-stress-test

Conversation

@drbenvincent

Copy link
Copy Markdown
Collaborator

Closes #427, closes #428, closes #429, closes #430, closes #431, closes #432, closes #433. Implements Phase 1 of #88 (tracked separately there). Meta: #426.

Summary

MMM/funnel generative workflows as an integration stress-test for pathmc's panel + simulate APIs. One commit-checkpoint per issue, each reviewed in-loop; full run report with per-issue evidence is on the tracking issue.

simulate() stack

Panel expressiveness

  • Feature: multi-dimensional panel support (geo × brand and nested unit keys) #430 — multi-dimensional units: panel={"unit": ["geo", "brand"], "time": "week"} builds a composite unit key (North|Acme) and pools over its Cartesian product; rectangularity validated with actionable errors; public results keep the flat (chain, draw, unit) dims contract.
  • Feature: MMM-style structured hierarchical pooling across panel dimensions #431 — structured pooling: pooling={"by_var": {"tv": {"coefficient": ("geo",)}, "theta_tv": "none"}} pools individual predictors/transform params over requested dimension subsets (mu_{var}_{dim}/sigma_{var}_{dim} hyperpriors) or forces per-cell params. Flat pooling dict unchanged. Transform-leaf predictors rejected at parse time (pooling raw columns under transformed terms would silently zero the media signal).
  • Feature: panel scaling helpers for heterogeneous units (MMM-style normalization) #432 — scaling: new pathmc.Scaling (max / mean / fixed / divide-by-grid methods, per-panel-dimension groups, xarray/dict grids aligned by dim name). Estimation runs on scaled columns; simulate() inverts generated endogenous columns back to business units. predict()/do() stay in scaled units this phase (documented; escape hatch model.fitted_scaling).
  • Panel latent dynamics: AR latent states with lag syntax #88 Phase 1 — latent AR dynamics: latents feeding a lag() term get estimated init_{var} initial conditions seeding the scan carry; new PathModel.latent_trajectory(var) returns posterior latent states as xarray (chain, draw, time, unit).

Docs/tests

Verification

  • Baseline on main before the run: make test-fast = 1160 passed / 0 failed (tag pre-run/426-2026-08-24).
  • Final on this branch: make test = 1448 passed / 0 failed, coverage 90.35%; make check_lint clean (ruff + mypy); uv run great-docs build succeeds.
  • Per-issue evidence commands are in the progress comments on each sub-issue.

Review notes

Each checkpoint went through an inline review loop; three real defects were caught and fixed with regression tests (endo-column init leak in #427, transform-leaf pooling corruption in #431, _construction regression breaking refute_placebo in #432). Taste-level decisions are recorded as ## Decision (2026-08-24) sections on the sub-issues; the ones worth a human eye are listed in the run report on #426 — notably that predict()/do() outputs remain in scaled units under #432.

Known follow-ups (not addressed here): pre-existing do(set={latent_var}) KeyError for stochastic-latent panels; literal transform args becoming RVs named '0.7'; temporal-dynamics predicate duplicated across modules.

…nally (#434)

Phase 1: draw correlated residuals per block from chol_{block} (packed
LKJCholeskyCov form); block members propagate as mu deterministics to
descendants, mirroring estimation. Phase 2: hsgp basis coefficients are
ordinary free RVs handled by the existing pm.do/draw path. Both
NotImplementedError guards removed; round-trip tests added.

Reviewed (round 1): APPROVE; taste notes applied (docstring propagation
semantics, zero-noise smooth check in hsgp test).

Verification: uv run pytest tests/test_simulate_data.py
tests/test_hsgp_compile.py tests/test_residual_cov.py tests/test_hsgp_do.py -q
=> 45 passed.
Enumerates every free-RV name simulate() requires, with concrete shapes
(unit-indexed hierarchical params, packed chol_{block}, hsgp basis
weights) via the same zero-filled placeholder compile simulate() uses.
Clear ValueError when lag()-bearing specs omit panel=. Documented in
data_simulation.qmd.

Reviewed (round 1): APPROVE; taste notes applied (IntoFrame annotation,
hsgp docstring wording). Literal-valued transform args becoming RVs
flagged upstream as follow-up.

Verification: pytest tests/test_simulate_template.py
tests/test_public_api.py tests/test_simulate_data.py -q => green.
Cross-sectional draw path reuses the name new_columns after the scan
branch; rename to new_columns_xs. No behavior change (51 tests green:
simulate_data, residual_cov, stochastic_latent).
pooling={'intercept': True, 'slopes': [...], 'by_var': {'tv':
{'coefficient': ('geo',)}, 'theta_tv': 'none'}} pools individual
predictors/transform params over requested panel-dimension subsets
(hyperpriors mu_{var}_{dim}/sigma_{var}_{dim}); 'none' forces per-cell
parameters. Flat pooling dict unchanged; legacy RV names untouched.
Transform-leaf predictors are rejected at parse time (pooling a raw
column under a transformed term would silently zero the transformed
signal); multi-equation and naming collisions validated early;
dim-less user prior overrides are re-dims'd to preserve per-cell
semantics in both compile paths. _warn_partial_pooling_intercept now
conditional per equation. Documented in 16-panel-data.qmd.

Reviewed: round 1 REQUEST_CHANGES (transform-leaf corruption, two
majors) -> fixed; round 2 APPROVE, no findings.

Verification: test_pooling_dims 22 passed incl. slow MCMC recovery
(mu_tv_geo [1.478,2.990] vs truth [1.5,3.0]); pooled suites 95 passed;
simulate/template 36 passed; ruff + mypy clean.
mmm_transform_data fixture, the adstock-decay recovery DGP, and
mmm_awareness.qmd Part 1 now generate ground truth via simulate()
(panel=, pooling=, latent AR(1) with estimated init). New fast
TestSimulateFunnelDGP: x2~x1 / x3~x1 / x4~x2+x3 / y ~
logistic_saturation(adstock(x4)) with zero-noise NumPy-reference
equality derived from the params dict. data_simulation.qmd callout
cross-links the MMM examples; freeze caches refreshed for both edited
pages.

Reviewed: round 1 REQUEST_CHANGES (orphaned old-cell tail in
mmm_awareness.qmd + stale _freeze) -> fixed; round 2 APPROVE.

Verification: test_v03_smoke + test_simulate_data + test_panel_latent
green incl. slow classes.
Resolved conflicts across panel.py, compile.py, priors.py, _model.py,
__init__.py, transforms.py, and docs/user_guide/18-panel-data.qmd.

Key integration decisions:
- build_panel_info keeps main's signature and validation (#388: null /
  duplicate / balance checks, require_rectangular flag); multi-D
  composite unit keys from #430 materialize via the new
  attach_composite_unit() helper called by model()/simulate()/
  simulate_params_template().
- by_var pooling machinery (#431) re-applied on top of main's MuSpec
  intercept rework (#417): pooled predictors excluded from flat beta in
  both compile paths; contributions added to mu_gen AND mu_est.
- Main's auto-drop of redundant partial-pooling intercepts supersedes
  the per-equation warning added for #431 (deleted upstream with its
  caller); tests updated to the new semantics.
- transforms.py: kept main's normalize-aware apply_pymc panel branch
  (supersedes #428's deletion of _apply_pymc_panel).
- Docs: our sections merged into renamed 18-panel-data.qmd with main's
  new link paths; freeze cache refreshed at main's new _freeze/docs/
  layout.

Verification on merged tree: 1862 passed / 2 skipped, coverage 90.90%.
@drbenvincent
drbenvincent marked this pull request as ready for review August 24, 2026 19:36
@daimon-pymclabs

Copy link
Copy Markdown
Collaborator

Adversarial review

Reviewed at branch head 21da9df (merge-base 61275b3), diff restricted to pathmc/ (1841 added / 120 removed). Every finding below was executed against the branch in a clean pathmc-dev environment (editable install of this commit) rather than reasoned from the diff alone — repro snippets and actual observed output are inline. make test being green is consistent with all of these: each one sits in a hole in the new test matrix, and I name the hole in every case.

Ordered by severity. Six blockers, then four consistency issues, then the test-matrix gaps.


Blockers

B1. simulate() silently drops every edge out of a residual-block member

Where: pathmc/_model.py (removal of the spec.residual_covs NotImplementedError), interacting with _compile_residual_block + _make_cross_sectional_resolver(prefer_observed_block_members=True) in pathmc/compile.py.

What happens. The PR removes this guard:

if spec.residual_covs:
    raise NotImplementedError("simulate() does not yet support residual covariances (~~). ...")

and adds a numpy post-hoc Cholesky draw. But block members are compiled as endogenous_rvs[var] = mu_det (a Deterministic), and any downstream equation that references a block member is built with prefer_observed_block_members=True, which strips block members from active_endogenous so the resolver falls through to the data column. Inside simulate() that data column is the zero-fill:

data_sim = nw_data.with_columns([nw.lit(0.0).alias(var) for var in endogenous_lhs])

_emit_free_rv(var, mu_est, ...) uses mu_est. So a descendant of a block member is generated from a mean structure in which its parent is identically zero. The correlated noise is then pasted onto the block members afterwards, by which point the descendants have already been drawn.

Observed (this branch):

out = pathmc.simulate(
    "Y1 ~ X\nY2 ~ X\nZ ~ Y1\nY1 ~~ Y2",
    data=pd.DataFrame({"X": rng.normal(size=4000)}),
    params={"beta_Y1": [0.0, 0.5], "beta_Y2": [0.0, -0.5],
            "beta_Z": [0.0, 3.0], "sigma_Z": 0.1,
            "chol_Y1_Y2": [1.0, 0.8, 0.6]},
    random_seed=42,
)
np.polyfit(out["Y1"], out["Z"], 1)[0]   # -> -0.0012   (true value: 3.0)
np.corrcoef(out["Y1"], out["Z"])[0, 1]  # -> -0.0136
np.var(out["Z"])                        # -> 0.01  == sigma_Z**2, i.e. pure noise

No warning, no error. The returned frame looks completely normal.

Why this is the worst one. simulate() exists to produce ground truth for simulate-and-recover. Silently generating data from a different DAG than the one the user wrote is the single failure mode this function must never have — a recovery test built on it will "pass" while validating nothing, or will report a real coefficient as unrecoverable and send someone hunting a phantom sampler bug. It is strictly worse than the NotImplementedError this PR deletes, because the old behaviour was loud.

Also note the new docstring states the opposite of what the code does:

block members' simulated columns share the correlated residuals, and descendant equations see their mean structure (mu_{var}), not the noisy realized values — mirroring estimation

Two things wrong. (a) Descendants see neither mu_{var} nor the realized values; they see the zero-filled pm.Data column. (b) Even the intended behaviour would not "mirror estimation": estimation is precisely the prefer_observed_block_members=True path, where descendants read the observed, noisy column. Generating from mu and estimating against realized values is textbook attenuation — the downstream coefficient would be biased even after the zero-fill bug is fixed.

Scaffolded fix. Two options; I'd ship (a) in this PR and open an issue for (b).

(a) Restore a scoped guard — smallest correct change. Keep the ~~ support that is actually exercised (blocks with no downstream consumers, which is what test_residual_cov_round_trip covers) and reject the broken shape explicitly:

# pathmc/_model.py, in simulate(), after block_var_set, blocks = _identify_residual_blocks(spec)
if block_var_set:
    consumers = {
        reg.lhs
        for reg in spec.regressions
        if reg.lhs not in block_var_set
        and block_var_set & {v for t in reg.terms for v in _term_base_vars(t)}
    }
    if consumers:
        raise NotImplementedError(
            f"simulate() cannot yet generate {sorted(consumers)}: they are "
            f"downstream of residual-covariance block member(s) "
            f"{sorted(block_var_set & {v for reg in spec.regressions if reg.lhs in consumers for t in reg.terms for v in _term_base_vars(t)})}. "
            "The generative graph wires block members through their mean "
            "structure, so the realized correlated residuals would not reach "
            "the descendant. Simulate the block in one call, then simulate "
            "descendants in a second call using the returned columns."
        )

Add a regression test asserting the raise, plus the positive case that already exists.

(b) Real fix — make block realizations first-class in the graph. The post-hoc numpy Cholesky is what forces the ordering problem: the correlated draw happens outside the pytensor graph, so nothing downstream can consume it. Move it inside. In _compile_residual_block, alongside the observed MvNormal, register a generative realization:

# after structure.emit(...)
chol = <the packed chol RV emitted by LKJResidual>
eps = pm.Normal(f"_resid_eps_{'_'.join(block_sorted)}", 0.0, 1.0,
                shape=(len(data), len(block_sorted)))
realized = pt.stack([mu_dict[v] for v in block_sorted], axis=1) + eps @ chol.T
for j, var in enumerate(block_sorted):
    generative_rvs[var] = realized[:, j]          # NOT endogenous_rvs

Then have _make_cross_sectional_resolver take a third mode ("generative") that prefers generative_rvs for block members, and build mu_gen with it while mu_est keeps today's observed-data wiring. simulate() then drops its numpy block entirely and just draws fixed_model[var] for block members like any other endogenous variable, with chol_{block} clamped via the existing pm.do path — which also gets pm.do() interventions on block members working for free. This is a bigger change; it deserves its own issue and its own recovery test (fit the model back on simulated data and check beta_Z recovers 3.0 inside the HDI).


B2. Multi-dimensional by_var coefficient pooling — the headline #431 feature — does not compile

Where: pathmc/compile.py, _compile_by_var_coefficients and the equivalent block in _compile_scan_panel.

What happens. _build_cell_group_index builds group_idx as a raveled index into the Cartesian product of the dim levels:

shape = tuple(len(levels[d]) for d in dims)
...
group_idx[ci] = np.ravel_multi_index(idxs, shape)

but the hyperprior it indexes is created with the dims unraveled:

mu_hp = _ensure_dims(priors[f"mu_{name}_{key}"], entry["dims"]).create_variable(f"mu_{name}_{key}")
beta  = pm.Normal(f"beta_{name}", mu=mu_hp[entry["group_idx"]], sigma=sigma_hp, dims="unit")

For dims=("geo","brand"), mu_hp is a (n_geo, n_brand) tensor and mu_hp[group_idx] is basic indexing on axis 0 only. The flat index is meaningless there, and it either silently selects the wrong row or blows up on shape.

Observed (3 geos x 2 brands x 12 weeks):

pathmc.model("y ~ tv", data=pdf,
             panel={"unit": ["geo", "brand"], "time": "week"},
             pooling={"by_var": {"tv": {"coefficient": ("geo", "brand")}}})
ValueError: Size length is incompatible with batched dimensions of parameter 0
AdvancedSubtensor{idx_list=(0,)}.0: len(size) = 1,
len(batched dims ...) = 2. Size must be None or have length >= 2

The identical spec with ("geo",) compiles fine (mu_tv_geo shape (3,)).

Why this is a blocker. The PR body sells #430 and #431 together — panel={"unit": ["geo","brand"]} plus "pools individual predictors/transform params over requested dimension subsets". The subset case (("geo",) out of ("geo","brand")) works; the full case (("geo","brand")) is the natural first thing a user types and it dies on a pytensor internals error that gives them no idea what they did wrong. _parse_by_var_pooling explicitly validates and accepts a tuple of dims, and default_priors explicitly sets dims=entry["dims"] — so the API advertises rank-N and the compiler only implements rank-1.

Root cause of it surviving review: tests/test_pooling_dims.py is 447 lines and every single coefficient entry in it is a one-element tuple (("geo",), ("channel",)). There is no rank-2 test anywhere in the file.

Scaffolded fix. Pick one and make the parser enforce it.

Preferred — flatten the hyperprior to match the flat index. Keep group_idx raveled and give the hyperprior a single derived coord:

# in _build_cell_group_index, also return the raveled coord
flat_levels = ["|".join(map(str, combo)) for combo in itertools.product(*(levels[d] for d in dims))]
return group_idx, levels, flat_levels

# in compile, register the derived coord and use it
coords[entry["key"]] = flat_levels          # e.g. coords["geo_brand"] = ["N|A", "N|B", ...]
mu_hp = _ensure_dims(priors[f"mu_{name}_{key}"], (entry["key"],)).create_variable(...)

and correspondingly in priors.default_priors / introspect.build_priors:

priors[f"mu_{name}_{key}"] = Prior("Normal", mu=0, sigma=10, dims=(entry["key"],))

This makes mu_{var}_{geo_brand} a (n_geo * n_brand,) vector indexed by the flat group_idx, which is exactly what the current indexing expression wants. Downside: the posterior is a flat vector, so slicing by geo means string-splitting the coord.

Alternative — keep the rank-N tensor and index it properly. Store the per-dim index arrays instead of the raveled one and use advanced indexing:

entry["dim_idx"] = tuple(np.array([level_idx[d][label_to_vals[lab][k]] for lab in panel_info.unit_labels])
                         for k, d in enumerate(dims))
...
mu_cell = mu_hp[entry["dim_idx"]]      # tuple of index arrays -> (n_cells,)
beta = pm.Normal(f"beta_{name}", mu=mu_cell, sigma=sigma_hp, dims="unit")

This keeps the nicer (geo, brand)-shaped posterior. It is the better end state; it needs _build_cell_group_index to return per-dim arrays and both call sites (cross-sectional and scan) updated.

Either way, guard the gap now so nobody hits the pytensor error:

# _parse_by_var_pooling, after dims are validated
if len(dims) > 1:
    raise NotImplementedError(
        f"by_var coefficient pooling over multiple dims {dims} for '{name}' is "
        "not implemented yet (see #<new issue>). Pool over a single panel "
        "dimension, e.g. {'coefficient': ('geo',)}, or use \"none\" for a "
        "fully unpooled per-cell coefficient."
    )

Required test either way — the one that is missing:

def test_by_var_coefficient_two_dims_recovers_cell_slopes(multi_dim_panel):
    m = pathmc.model("y ~ tv", data=multi_dim_panel,
                     panel={"unit": ["geo", "brand"], "time": "week"},
                     pooling={"by_var": {"tv": {"coefficient": ("geo", "brand")}}})
    assert tuple(m.pymc_model["mu_tv_geo_brand"].shape.eval()) == (3, 2)   # or (6,)
    # and a simulate-and-recover asserting per-cell betas land in the HDI

B3. ScalingFactors are positionally indexed, so reuse on any other frame is silently wrong

Where: pathmc/scaling.py, ScalingFactors.factors: dict[str, np.ndarray] (one element per data row), transform() and inverse_transform_column().

What happens. Factors are stored as a per-row array aligned to the estimation frame's row order. transform() does a bare positional divide:

vals = df[col].to_numpy().astype(float)
new_cols.append(nw.new_series(col, vals / factor, backend=df.implementation))

There is nothing tying a factor to the unit it was computed for. The docstring of fitted_scaling actively invites reuse: "Pass this object back as scaling= to pathmc.simulate() to reuse the exact estimation-time scales." Any frame that is not byte-for-byte in the same row order silently gets the wrong divisors.

Observed (2 geos, tv = 10 for N, 1000 for S, method="max" per geo; correct output is all 1.0):

f = m.fitted_scaling
f.transform(nw.from_native(sdf.iloc[::-1].reset_index(drop=True), eager_only=True))
# tv -> [1e+02]*8 + [1e-02]*8      # off by 100x in both directions, no error

And a differently-sized frame gives a raw numpy message with no context:

f.transform(nw.from_native(sdf.head(4), eager_only=True))
# ValueError: operands could not be broadcast together with shapes (4,) (16,)

Why this is an issue. The 100x error is silent and directionally plausible — a user comparing simulated to observed magnitudes would see "roughly the right shape, wrong level" and reach for the model before suspecting the scaler. And the intended workflow is cross-frame: fit on the estimation panel, simulate on a counterfactual/extended frame. A row-count mismatch is the expected case there (e.g. simulating a longer horizon), and it hard-errors with operands could not be broadcast. The per-row representation is also O(n_rows) memory for what is O(n_units) information.

Scaffolded fix. Store the factors keyed by unit, materialize per-row on demand.

@dataclass
class ScalingFactors:
    #: column -> (dims, {unit_key_tuple: divisor}). dims=() means a single global scale.
    factors: dict[str, tuple[tuple[str, ...], dict[tuple[str, ...], float]]] = field(default_factory=dict)

    def _per_row(self, df: nw.DataFrame, column: str) -> np.ndarray:
        dims, table = self.factors[column]
        keys = _row_keys(df, dims)
        missing = sorted({k for k in keys if k not in table})
        if missing:
            raise KeyError(
                f"Scaling factors for column {column!r} have no entry for "
                f"unit(s) {missing[:5]}. The factors were fitted on units "
                f"{sorted(table)[:5]}...; pass a frame whose {dims} values are "
                "a subset of the estimation-time units, or refit the scaling."
            )
        return np.array([table[k] for k in keys])

    def transform(self, df):
        return df.with_columns([
            nw.new_series(col, df[col].to_numpy().astype(float) / self._per_row(df, col),
                          backend=df.implementation)
            for col in self.factors if col in df.columns
        ])

    def inverse_transform_column(self, values, column, df):
        return np.asarray(values, float) * self._per_row(df, column)

_fit_role already computes exactly this table internally (_group_stat's stats dict, _grid_factors' table) and then throws the keys away by expanding to rows — so this is mostly deleting the expansion step. _invert_generated_columns in _model.py needs the frame threaded through, which it already has in scope.

Add tests: same factors applied to (i) a reversed frame, (ii) a subset frame, (iii) a frame with an unseen unit label — asserting correct values for i/ii and an actionable KeyError for iii.


B4. Composite unit keys silently merge distinct units

Where: pathmc/panel.py, _composite_unit_column.

What happens. The composite key is an unescaped "|" join with no validation that "|" is absent from the source values:

keys = df[unit_columns[0]].cast(nw.String)
for col in unit_columns[1:]:
    keys = keys + "|" + df[col].cast(nw.String)

("A|B", "C") and ("A", "B|C") both produce "A|B|C".

Observed:

# non-rectangular / cross-sectional panel -> SILENT MERGE
build_panel_info(df, {"unit": ["geo", "brand"], "time": "week"}, require_rectangular=False)
# .unit_labels -> ['A|B|C']      # two real units collapsed into one
# rectangular panel -> misleading error that blames the user's data
ValueError: Panel data has duplicate (unit, time) rows, which cannot be reshaped
into a rectangular panel: ('A|B|C', 0), ... Each combination of 'geo|brand' and
'week' must appear exactly once.

Why this is an issue. Pipe characters are not exotic in real panel keys — concatenated campaign/creative IDs, hierarchical SKU codes, scraped category paths. In the cross-sectional case two units are pooled into one and every per-unit parameter for them is wrong, with no signal at all. In the scan case the error message actively misdirects: the user's (geo, brand, week) rows are unique; it is pathmc's own derived key that collided, and nothing in the message says so.

Related: _as_grid in scaling.py inverts this join (key.split("|") when len(dims) > 1), so the same ambiguity corrupts scaling-grid lookups.

Scaffolded fix. Validate eagerly in build_panel_info, before the composite is built:

_SEPARATOR = "|"   # NB: this constant already exists in panel.py and is currently unused — wire it up

if len(unit_columns) > 1:
    for col in unit_columns:
        vals = df[col].cast(nw.String).unique().to_list()
        bad = sorted({v for v in vals if v is not None and _SEPARATOR in v})
        if bad:
            raise ValueError(
                f"Panel unit column {col!r} contains the composite-key separator "
                f"{_SEPARATOR!r} in value(s) {bad[:5]}. Multi-dimensional panels "
                f"join panel['unit'] columns with {_SEPARATOR!r}, so these values "
                "would make distinct units indistinguishable. Recode the column "
                "(e.g. replace the separator) before passing it as a unit column."
            )

Use _SEPARATOR everywhere instead of the three hardcoded "|" literals in panel.py and the two in scaling.py, so the separator can be changed in one place if this ever needs to be configurable.

Test: assert the raise for the colliding fixture above, and assert a non-colliding multi-dim panel still yields the expected unit_labels.


B5. ~42 lines of dead scan code that documents a fix which is not applied

Where: pathmc/compile.py ~L2386–2430, the lagged_exog_sequences block.

What happens. The block builds a dict of pre-computed lagged exogenous tensors, under a long comment explaining a pytensor scan-merge bug (pymc-devs/pytensor#2252, pathmc #316/#333) and stating:

Fix: build the lagged tensor directly from the existing pm.Data nodes ... and pass it as a plain scan sequence rather than carry state. This eliminates the trivial-echo carry that triggered the merge bug.

lagged_exog_sequences is never read:

$ grep -n lagged_exog_sequences pathmc/compile.py
2406:        lagged_exog_sequences: dict[str, Any] = {}
2412:                lagged_exog_sequences[base] = pt.concatenate(
2427:                lagged_exog_sequences[base] = pt.concatenate(

The sequences list passed to pm.scan does not include it, and outputs_info still contains [_init_carry(init_exog_lag[k]) for k in exog_lag_bases] — i.e. the carry path the comment says was eliminated is still the live path.

Why this is an issue. This is worse than ordinary dead code. It is 42 lines of confident, specific, cross-referenced documentation asserting that a known-nasty correctness bug is worked around here. The next person to touch the scan will read it, believe the exog-lag path is sequence-based, and reason incorrectly about the carry semantics. If the workaround is still needed, the PR ships a silent regression against #316; if it is not needed, the comment is a landmine. Neither is acceptable to merge as-is.

Scaffolded fix. Decide which, and say so in the code:

If the workaround is still required, wire it up:

sequences = (
    [exog_data_nodes[k] for k in exog_keys]
    + [lagged_exog_sequences[k] for k in exog_lag_bases]     # <- add
    + [observed_carry_nodes[k] for k in observed_carry_vars]
    ...
)
outputs_info = (
    [ ... endo ... ]
    + [_init_carry(init_adstock[k]) for k in adstock_keys]
    # - [_init_carry(init_exog_lag[k]) for k in exog_lag_bases]   <- remove
    + [None for _ in carry_mu_vars]
)

and update step_fn's signature/unpacking and n_seq accordingly, plus the regression test from #316 (fit a lag model, assert the lag coefficient's posterior is not concentrated at zero).

If it is no longer required (2252 fixed upstream, or the carry shape changed), delete the block and replace the comment with one line recording the resolution:

# Exog lags use scan carry state. The pytensor scan-merge workaround
# (pytensor#2252 / pathmc #333) was removed on <date> after <evidence>.

B6. n_steps guard removed; the comment justifying its replacement describes an unreachable branch

Where: pathmc/compile.py, the pm.scan(...) call.

-            n_steps=n_times,
+            # Pure latent dynamics (e.g. ``awareness ~ lag(awareness)``
+            # with no exogenous columns and a deterministic latent)
+            # leave the sequence list empty; scan then needs an explicit
+            # step count to know how long the recursion runs.
+            n_steps=n_times if not sequences else None,

Why this is an issue. sequences is never empty. It unconditionally ends with [do_intervene_nodes[k] for k in endo_keys], and do_intervene_nodes is built for every k in endo_keys = list(endogenous_order) — which is non-empty for any model that reaches _compile_scan_panel. So the not sequences branch is dead, the comment describes a case that cannot occur, and the actual effect of this hunk is unconditionally changing n_steps=n_times to n_steps=None.

That is a real behavioural change: scan now infers its trip count from the shortest sequence instead of being explicitly bounded. Today every sequence is (n_times, n_units) so it is equivalent — but the explicit bound was the thing that would catch a future length mismatch (a pm.set_data with a different horizon, a sequence built one row short) as a loud scan error instead of a silently truncated recursion. The PR trades a guard for nothing and labels the trade with an incorrect rationale.

Scaffolded fix. pm.scan accepts sequences and n_steps together (it truncates to n_steps). Restore the unconditional bound and drop the comment:

n_steps=n_times,

If the pure-latent case genuinely was hitting an error, that means sequences was empty at that moment — which would contradict the code above, so please attach the failing repro to the issue before re-introducing the conditional. If the conditional is kept for a reason I'm not seeing, assert the invariant so the dead branch is visible:

assert sequences, "scan sequences always include do_intervene nodes"

Consistency / API issues

C1. init_{var} priors are registered for models that never use them

priors.default_priors and introspect.build_priors both compute:

lag_base_vars = {term.lag_of for reg in spec.regressions for term in reg.terms if term.lag_of is not None}
...
if reg.lhs in lag_base_vars:
    priors[f"init_{reg.lhs}"] = Prior("Normal", mu=0, sigma=1)

with no panel gate — while the compiler only emits init_{var} inside _compile_scan_panel, for sorted(set(latent) & set(endo_lag_bases)). Confirmed on this branch:

default_priors(parse_spec("a ~ lag(a) + x\ny ~ a"), latent={"a"}, families={"a": "latent_normal"})
# keys: ['beta_a', 'beta_y', 'init_a', 'sigma_a', 'sigma_y']    <- init_a is never consumed

The comment directly above the code says "exist only for latent variables that feed a lag() term in scan-compiled panel models" — so the intent is documented and the implementation does not match it. Users calling .equations() on a non-panel spec see a phantom parameter, and set_priors("init_a", ...) is a silent no-op.

default_priors already receives panel_info. Gate on it:

is_scan_panel = panel_info is not None and _has_temporal_deps(spec, graph_info)
lag_base_vars = ({...}) if is_scan_panel else set()

introspect.build_priors needs the same signature/gate; today it duplicates the predicate a third time. The PR's own "known follow-ups" note flags "temporal-dynamics predicate duplicated across modules" — this is a good moment to hoist it into one compile._is_scan_panel(spec, graph_info, panel_info) helper and call it from all three sites, because the duplication has already produced a divergence.

C2. init_{var} gets dims only when the user overrides the prior

if priors and f"init_{var}" in priors:
    latent_init_rvs[var] = _ensure_dims(priors[f"init_{var}"], ("unit",)).create_variable(f"init_{var}")
else:
    latent_init_rvs[var] = pm.Normal(f"init_{var}", mu=0, sigma=1, shape=(n_units,))

The else branch uses shape= rather than dims="unit", so the posterior variable would land with an anonymous init_{var}_dim_0 instead of the unit coord. In practice default_priors always registers init_{var}, so the else is dead — which means the branch exists only to introduce an inconsistency if it ever becomes live. Collapse it:

latent_init_rvs[var] = _ensure_dims(priors[f"init_{var}"], ("unit",)).create_variable(f"init_{var}")

and let a KeyError here be the signal that prior registration and compilation have diverged (C1's real cause).

C3. init_{var} ~ Normal(0, 1) is not scale-free, on a variable with no data scale

The default initial condition for a latent is fixed at N(0, 1). Latents have no observed column, so their scale is set entirely by the priors on the rest of the recursion. For an awareness-style latent whose stationary level is O(10) or O(100), an N(0,1) init is a strongly informative prior pinning t=0 near zero, and the recursion takes 1/(1-decay)-ish steps to escape it — biasing exactly the early-period decomposition MMM users care about.

At minimum this needs a docstring warning at latent_trajectory() and in 18-panel-data.qmd telling users to override init_{var} when their latent is not O(1). Better: make it relative to the latent's own scale, e.g. default to Normal(0, sigma_{var} / sqrt(1 - decay**2)) for a stochastic latent with a single lag term (the stationary sd), falling back to N(0,1) when the recursion is not a simple AR(1). Worth a sensitivity check in the recovery test: fit the same simulated data with init prior sd 1 vs 10 and assert the decay posterior does not move.

C4. fixed / divide scaling accept negative and non-finite divisors

_fit_role rejects zero, and _group_stat requires max/mean scales to be positive and finite — but fixed and grid-based divisors are only checked against zero.

pathmc.model("y ~ tv", data=sdf, panel=sp,
             scaling=pathmc.Scaling(channel={"method": "fixed", "value": -2.0}))
# accepted; factors -> [-2., -2., -2.]

A negative divisor sign-flips the column, which flips the sign of every coefficient on it relative to the user's mental model — and the Scaling docstring says "a positive constant", so the check is simply missing rather than a deliberate allowance. NaN/inf in a divide grid is worse: it silently NaNs the column and the model fails later at sampling with an unrelated message.

Consolidate the validation the way _group_stat already does it:

def _require_positive_finite(arr: np.ndarray, *, role: str, method: str, where: str) -> None:
    bad = ~(np.isfinite(arr) & (arr > 0))
    if bad.any():
        raise ValueError(
            f"Scaling.{role} method {method!r} produced non-positive or "
            f"non-finite divisor(s) {np.unique(arr[bad])[:5].tolist()} for {where}. "
            "Every scale factor must be positive and finite."
        )

and call it at the end of all three _fit_role branches. (If negative divisors are ever legitimately wanted, that should be an explicit opt-in key, not the absence of a check.)

C5. Small ones

  • _SEPARATOR is dead. Defined in panel.py:58, never referenced; the separator is hardcoded as "|" in three places in the same file. Fix as part of B4.
  • _emit_transform_priors(transform_map=...) is unused — the parameter is accepted and ignored (pre-existing, but the signature was touched here to add existing=, so it's a free cleanup).
  • transform_param_rvs is overwritten with a row-expanded tensor in the cross-sectional path (transform_param_rvs[pname] = transform_param_rvs[pname][unit_idx]), and the two loops guard on different conditions — pname in priors for creation, pname in transform_param_rvs for expansion. If a transform param ever lacks a prior entry, the generic emitter creates it as a scalar and the second loop indexes a scalar with unit_idx. Make the second loop mirror the first (if entry["kind"] == "none_transform" and pname in priors) or track the names created by the first loop in a local set.
  • Silent no-op when unit_idx is None. if pooled_by_lhs.get(var) and unit_idx is not None: — the flat beta slot has already been zeroed by _exclude_pooled_from_flat_beta at that point, so if this guard ever fails the predictor vanishes from the model entirely with no error. Currently unreachable (_parse_by_var_pooling raises when panel_info is None), but make it explicit: assert unit_idx is not None, "by_var pooling requires a unit index".
  • Duplicate coord. For a single-column panel, coords ends up with both unit and the dim name carrying identical levels ({'unit': ('E','N','S'), 'geo': ('E','N','S')}). Harmless, but it doubles the coord in every idata; consider reusing unit when dims == panel_info.unit_columns.
  • _exclude_pooled_from_flat_beta docstring says "Mutates mu_specs in place" — it replaces the dict values with dataclasses.replace copies. The dict is mutated; the MuSpec objects are not. Worth being precise, since a reader relying on in-place slot mutation would write a subtly wrong follow-up.
  • simulate_params_template does not do what its docstring says. It claims to compile "the same zero-filled placeholder generative model that simulate() builds internally", but simulate() overwrites all endogenous columns with zeros while the template only fills missing ones, and the template skips scaling_factors.transform() entirely. Shapes are unaffected today, so nothing is broken — but the stated invariant is what makes the template trustworthy. Either factor the shared frame prep into one _prepare_simulation_frame(spec, nw_data, panel_info, scaling_factors) used by both, or soften the docstring to state exactly which parts are shared.
  • _fit_role recomputes _row_keys(df, dims) inside the per-column loop for max/meanO(n_cols * n_rows) string work for a value that does not vary. Hoist it above the loop.

Test-matrix gaps

Each blocker above corresponds to a specific hole. make test = 1448 passed is not evidence against them; it is evidence that the matrix does not reach them.

# Missing case Nearest existing test
B1 ~~ block member as a predictor of a downstream equation in simulate() test_residual_cov_round_trip — block members are terminal nodes
B2 {"coefficient": (d1, d2)} with two or more dims all 11 coefficient entries in test_pooling_dims.py are 1-tuples
B3 ScalingFactors reused on a reordered / subset / extended frame test_scaling.py only round-trips within the fitting frame
B4 unit column values containing the `" "` separator
B5/B6 scan structure assertions (which tensors are sequences vs carries; n_steps bound) none — scan wiring is only tested end-to-end via recovery
C4 negative / NaN / inf scaling divisors zero is tested; sign and finiteness are not

The pattern worth naming: every one of these is a silent-wrong-answer path, not a crash path, and the recovery-style tests this PR adds are structurally blind to them because they assert "posterior contains truth" on a DGP the same code generated. A simulate() bug that corrupts the DGP corrupts the assertion identically and the test still passes. For B1 and B3 specifically, please add at least one test that checks a generated column against a hand-written NumPy reference rather than against a pathmc-generated one — #428 already establishes exactly this pattern for transforms and it is the right pattern here too.


What I'd merge

The simulate() panel/transform work (#427/#428), simulate_params_template (#429), multi-dim panel keys (#430) modulo B4, single-dim by_var pooling (#431 — verified working: per-cell slopes 1.0/5.0/20.0 recovered to 3 decimal places through simulate()), and the latent AR work (#88) all look solid and are well tested.

Blocking merge: B1 (restore a scoped guard at minimum), B2 (fix or explicitly reject rank-N dims), B3 (unit-keyed factors), B4 (separator validation), B5 (resolve the dead scan block one way or the other), B6 (restore n_steps).

Happy to pair on any of these, or to split B1(b) and B2's preferred fix into their own issues if you'd rather land the guards here and do the real fixes next.

Verified against 21da9df in a clean pathmc-dev environment; all repro snippets above were executed, and the quoted outputs are actual, not illustrative.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

2 participants