feat(mmm): add ControlEffect MuEffect to handle control variables via protocol#2568
Draft
williambdean wants to merge 2 commits into
Draft
feat(mmm): add ControlEffect MuEffect to handle control variables via protocol#2568williambdean wants to merge 2 commits into
williambdean wants to merge 2 commits into
Conversation
… protocol Introduces ControlEffect, a new MuEffect subclass that routes control-variable logic through the MuEffect protocol (create_data / create_effect / set_data) instead of inline build_model code. - MuEffect ABC: create_data now accepts optional X: xr.Dataset | None so effects that need external data (e.g. control columns) can consume it while existing effects safely ignore it - ControlEffect: registers control_data as a pmd.Data shared variable, creates gamma_control from a configurable Prior, and exposes a control_contribution Deterministic; supports panel dims transparently - field_validator mirrors Transformation.with_default_prior_dims: dims=None on the supplied Prior is defaulted to 'control'; explicit dims are untouched - MMM.__init__: auto-inserts ControlEffect(prior=model_config['gamma_control']) at position 0 of mu_effects when control_columns are present - MMM.build_model: removed inline control block; now delegates to mu_effects - MMM._set_xarray_data: removed hardcoded control pm.set_data block - MMM.build_from_idata: backward-compat guard re-injects ControlEffect for models saved before this change - Serialization: ControlEffect registered in the serialization registry; save/load round-trip preserves prior and date_dim_name
Covers the full ControlEffect surface: - create_data registers shared variable and coord; is idempotent - create_data is a no-op when X is None or lacks '_control' key - create_effect produces a 'control_contribution' Deterministic with 'control' summed out; raises RuntimeError if create_data was skipped - set_data updates shared variable; is a no-op for None or missing key - to_dict / from_dict round-trip preserves prior and date_dim_name - field_validator: Prior(dims=None) defaults to dims='control'; explicitly supplied dims are left untouched - Parametrized scalar and panel-geo multidimensional cases
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
1 participant
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.
Context
This is exploratory work on how to allow MuEffects to consume additional external data (e.g. control columns from the training
XDataFrame). Previously, control-variable logic lived as inline code insidebuild_modeland_set_xarray_data. This PR routes it through theMuEffectprotocol instead, usingControlEffectas the first concrete example.The key design question being explored: how should MuEffects that need external data (beyond model coordinates) receive it? The solution here is to pass an optional
X: xr.Dataset | Nonetocreate_data, which effects that don't need external data simply ignore.Changes
MuEffectABCcreate_data(mmm, X=None)— optionalXso effects that need external data can consume it; existing effects ignore it safelyset_data(mmm, model, X: xr.Dataset | None)— type annotation correctedControlEffect(MuEffect)control_dataas apmd.Datashared variablegamma_controlfrom a configurablePriorcontrol_contributionDeterministic with"control"dim summed outfield_validatormirrorsTransformation.with_default_prior_dims:dims=Nonedefaults to"control"; explicit dims are untouchedMMM__init__: auto-insertsControlEffect(prior=model_config["gamma_control"])at position 0 ofmu_effectswhencontrol_columnsare presentbuild_model: inline control block removed; delegated tomu_effects_set_xarray_data: hardcoded controlpm.set_datablock removedbuild_from_idata: backward-compat guard re-injectsControlEffectfor models saved before this changeWhat's not in scope
X(onlyControlEffectuses it today)control_columnsparam unchanged)mmm.dimsTesting
TestControlEffectunit tests covering create/effect/set_data, idempotency, serialization round-trip, prior validator, and scalar + panel-geo multidimensional cases📚 Documentation preview 📚: https://pymc-marketing--2568.org.readthedocs.build/en/2568/