Support fixed/user-defined scaling factors for stable model refreshes (#2479)

* Support fixed/user-defined scaling factors for stable model refreshes (#2478)

Add method="fixed" to VariableScaling so users can pin scaling constants
that remain stable across production model refreshes, eliminating
scale-induced prior drift when data distributions change.

Made-with: Cursor

* Address review feedback: serialization, dict-valued scaling, key validation

- Fix legacy MMM serialization round-trip: use model_dump() in
  create_idata_attrs and pass value through in _deserialize_scaling so
  fixed scaling survives save/load.
- Support dict-valued fixed channel scaling in legacy MMM (per-channel
  constants keyed by column name); reject dict-valued target with a
  clear error since the 1-D MMM has a single target.
- Validate dict keys against coordinate labels in multidimensional MMM:
  report missing and unexpected keys explicitly instead of late KeyError
  or silent ignore.
- Add regression tests for all three fixes.

Made-with: Cursor

* Improve patch coverage: add tests, remove dead code

- Remove dead `scaling is None` branch in create_idata_attrs (scaling is
  always set in __init__)
- Replace isinstance guard with cast() for mypy narrowing
- Add test for dict-valued scaling with wrong remaining dims count
- Add tests for VariableScaling dims validation (date, duplicates)

Made-with: Cursor

* Restore legacy MMM scaling type check

This preserves the legacy TypeError contract for non-fixed scaling so the shard 5 edge-case test passes after the fixed-scaling refactor.

Made-with: Cursor

* Refactor VariableScaling into a two-class hierarchy

Split the single VariableScaling class into an abstract base with two
concrete subclasses (DataDerivedScaling and FixedScaling) to eliminate
the awkward optional value field and its cross-field validation, per
code review feedback.

- VariableScaling(ABC) keeps dims + _validate_dims
- DataDerivedScaling has method: Literal["max", "mean"]
- FixedScaling has value: float | dict[str, float] (non-optional)
- Scaling.to_dict/from_dict use the __type__ registry for polymorphic
  serialization; legacy format (no __type__) still loads correctly
- Updated legacy MMM and multidimensional MMM to use isinstance checks
  and the new concrete classes
- All 844 tests pass across test_scaling, test_mmm, test_multidimensional

Made-with: Cursor

* Extract _build_fixed_scale to eliminate multiple returns

Address review comment by refactoring _compute_scale_for_variable into
two methods: the main method now has a single return, dispatching to a
dedicated _build_fixed_scale helper for FixedScaling logic.

Made-with: Cursor

* Validate dict-valued FixedScaling keys at init, not build_model

Move the missing/extra key checks for dict-valued channel FixedScaling
from build_model() to MMM.__init__() in both legacy and multidimensional
implementations, so users get immediate feedback on key mismatches.

Made-with: Cursor

* feat(mmm): multidimensional fixed scaling with DataArray and init fix

- Validate dict channel keys at init only when remaining dim is channel
- Support FixedScaling.value as xarray.DataArray (broadcast to reduced grid)
- JSON-safe DataArray serialization; from_long_dataframe helper
- Clearer dict error when multiple remaining dims; target dict keys at build
- Tests: broadcast, full grids, 3D panel, misaligned coords, idata round-trip
- Doc examples for panel fixed scaling

Made-with: Cursor

* fix(test): add fastprogress to test extras for BlackJAX mlflow autolog

BlackJAX sampling uses jax.io_callback for progress bars, which requires
fastprogress. CI installs only .[test]; without it, test_autolog_pymc_model
fails with CpuCallback errors on Linux runners.

Made-with: Cursor

* revert: drop fastprogress from test extras (unrelated to fixed scaling PR)

Restore prior test dependency set; BlackJAX/mlflow env gap can be
addressed separately on main.

Made-with: Cursor

Author: drbenvincent

pymc_marketing/mmm/__init__.py

@@ -63,6 +63,12 @@
     preprocessing_method_X,
     preprocessing_method_y,
 )
+from pymc_marketing.mmm.scaling import (
+    DataDerivedScaling,
+    FixedScaling,
+    Scaling,
+    VariableScaling,
+)
 from pymc_marketing.mmm.sensitivity_analysis import SensitivityAnalysis
 from pymc_marketing.mmm.time_slice_cross_validation import (
     TimeSliceCrossValidationResult,
@@ -78,8 +84,10 @@
     "BaseValidateMMM",
     "BinomialAdstock",
     "CovFunc",
+    "DataDerivedScaling",
     "DelayedAdstock",
     "FancyLinearRegression",
+    "FixedScaling",
     "GeometricAdstock",
     "HSGPPeriodic",
     "HillSaturation",
@@ -99,12 +107,14 @@
     "PeriodicCovFunc",
     "RootSaturation",
     "SaturationTransformation",
+    "Scaling",
     "SensitivityAnalysis",
     "SoftPlusHSGP",
     "TanhSaturation",
     "TanhSaturationBaselined",
     "TimeSliceCrossValidationResult",
     "TimeSliceCrossValidator",
+    "VariableScaling",
     "WeeklyFourier",
     "WeibullCDFAdstock",
     "WeibullPDFAdstock",

pymc_marketing/mmm/mmm.py

@@ -17,7 +17,7 @@
 import logging
 import warnings
 from collections.abc import Sequence
-from typing import Annotated, Any, Literal
+from typing import Annotated, Any, Literal, cast
 
 import arviz as az
 import matplotlib.pyplot as plt
@@ -54,7 +54,13 @@
     scale_lift_measurements,
 )
 from pymc_marketing.mmm.preprocessing import MaxAbsScaleChannels, MaxAbsScaleTarget
-from pymc_marketing.mmm.scaling import Scaling, VariableScaling
+from pymc_marketing.mmm.scaling import (
+    DataDerivedScaling,
+    FixedScaling,
+    Scaling,
+    deserialize_variable_scaling,
+    validate_fixed_scaling_keys,
+)
 from pymc_marketing.mmm.tvp import create_time_varying_gp_multiplier, infer_time_index
 from pymc_marketing.mmm.utility import UtilityFunctionType, average_response
 from pymc_marketing.mmm.utils import (
@@ -205,22 +211,32 @@ def __init__(
         self.validate_data = validate_data
         self.adstock_first = adstock_first
 
-        # Initialize scaling configuration similar to multidimensional MMM
         if isinstance(scaling, dict):
             scaling = scaling.copy()
 
             if "channel" not in scaling:
-                scaling["channel"] = VariableScaling(method="max", dims=())
+                scaling["channel"] = DataDerivedScaling(method="max", dims=())
             if "target" not in scaling:
-                scaling["target"] = VariableScaling(method="max", dims=())
+                scaling["target"] = DataDerivedScaling(method="max", dims=())
 
             scaling = Scaling(**scaling)
 
         self.scaling: Scaling = scaling or Scaling(
-            target=VariableScaling(method="max", dims=()),
-            channel=VariableScaling(method="max", dims=()),
+            target=DataDerivedScaling(method="max", dims=()),
+            channel=DataDerivedScaling(method="max", dims=()),
         )
 
+        validate_fixed_scaling_keys(self.scaling.channel, channel_columns, "channel")
+
+        if isinstance(self.scaling.target, FixedScaling) and isinstance(
+            self.scaling.target.value, dict
+        ):
+            raise ValueError(
+                "Dict-valued fixed target scaling is not supported in the "
+                "legacy MMM (single target). Use a scalar value or switch "
+                "to the multidimensional MMM."
+            )
+
         model_config = model_config or {}
         model_config = parse_model_config(
             model_config,  # type: ignore
@@ -450,24 +466,59 @@ def _compute_scale_for_data(
 
     def _compute_scales(self) -> None:
         """Compute and save scaling factors for channels and target."""
-        # Get raw data
-        X_data = self.preprocessed_data["X"]
-        if not isinstance(X_data, pd.DataFrame):
-            raise TypeError("X data must be a DataFrame for scaling computation")
-
-        # Use pandas/numpy efficient operations - avoid redundant .values call
-        channel_data = X_data[self.channel_columns].to_numpy()
-        target_data = np.atleast_1d(np.asarray(self.preprocessed_data["y"]))
+        channel_scaling = self.scaling.channel
+        target_scaling = self.scaling.target
+
+        channel_scale: np.ndarray | float
+        if isinstance(channel_scaling, FixedScaling):
+            if isinstance(channel_scaling.value, dict):
+                channel_scale = np.array(
+                    [channel_scaling.value[c] for c in self.channel_columns],
+                    dtype=float,
+                )
+            elif isinstance(channel_scaling.value, DataArray):
+                raise ValueError(
+                    "DataArray-valued FixedScaling is not supported by the "
+                    "legacy MMM. Use a scalar or dict value, or switch to "
+                    "MultidimensionalMMM."
+                )
+            else:
+                n_channels = len(self.channel_columns)
+                channel_scale = np.full(n_channels, channel_scaling.value)
+        else:
+            X_data = self.preprocessed_data["X"]
+            if not isinstance(X_data, pd.DataFrame):
+                raise TypeError("X data must be a DataFrame for scaling computation")
 
-        # Compute scales based on scaling configuration
-        self.channel_scale = self._compute_scale_for_data(
-            channel_data, self.scaling.channel.method, axis=0
-        )
-        target_scale = self._compute_scale_for_data(
-            target_data, self.scaling.target.method, axis=None
-        )
-        # Ensure target_scale is a Python float (convert from numpy scalar if needed)
-        self.target_scale = float(target_scale)
+            X_data = cast(pd.DataFrame, X_data)
+            channel_data = X_data[self.channel_columns].to_numpy()
+            channel_scale = self._compute_scale_for_data(
+                channel_data, channel_scaling.method, axis=0
+            )
+        self.channel_scale = channel_scale
+
+        target_scale: float
+        if isinstance(target_scaling, FixedScaling):
+            if isinstance(target_scaling.value, DataArray):
+                raise ValueError(
+                    "DataArray-valued FixedScaling is not supported by the "
+                    "legacy MMM. Use a scalar or dict value, or switch to "
+                    "MultidimensionalMMM."
+                )
+            if not isinstance(target_scaling.value, (int, float)):
+                raise TypeError(
+                    f"Expected scalar FixedScaling value for target, "
+                    f"got {type(target_scaling.value).__name__}."
+                )
+            target_scale = float(target_scaling.value)
+        else:
+            target_data = np.atleast_1d(np.asarray(self.preprocessed_data["y"]))
+            target_scale = float(
+                self._compute_scale_for_data(
+                    target_data, target_scaling.method, axis=None
+                )
+            )
+        self.target_scale = target_scale
 
     def create_idata_attrs(self) -> dict[str, str]:
         """Create attributes for the inference data.
@@ -493,22 +544,9 @@ def create_idata_attrs(self) -> dict[str, str]:
         attrs["treatment_nodes"] = json.dumps(self.treatment_nodes)
         attrs["outcome_node"] = json.dumps(self.outcome_node)
 
-        # Serialize scaling configuration
-        if hasattr(self, "scaling") and self.scaling is not None:
-            attrs["scaling"] = json.dumps(
-                {
-                    "target": {
-                        "method": self.scaling.target.method,
-                        "dims": self.scaling.target.dims,
-                    },
-                    "channel": {
-                        "method": self.scaling.channel.method,
-                        "dims": self.scaling.channel.dims,
-                    },
-                }
-            )
-        else:
-            attrs["scaling"] = json.dumps(None)
+        from pymc_marketing.serialization import serialization as _serialization
+
+        attrs["scaling"] = json.dumps(_serialization.serialize(self.scaling))
 
         return attrs
 
@@ -1254,6 +1292,9 @@ def _data_setter(
     def _deserialize_scaling(cls, scaling_dict: dict | None) -> Scaling | None:
         """Deserialize scaling configuration from JSON.
 
+        Handles both new format (with ``__type__`` keys from the serialization
+        registry) and legacy format (flat dicts with ``method``/``dims``).
+
         Parameters
         ----------
         scaling_dict : dict | None
@@ -1267,15 +1308,14 @@ def _deserialize_scaling(cls, scaling_dict: dict | None) -> Scaling | None:
         if scaling_dict is None:
             return None
 
+        if "__type__" in scaling_dict:
+            from pymc_marketing.serialization import serialization as _serialization
+
+            return _serialization.deserialize(scaling_dict)
+
         return Scaling(
-            target=VariableScaling(
-                method=scaling_dict["target"]["method"],
-                dims=tuple(scaling_dict["target"]["dims"]),
-            ),
-            channel=VariableScaling(
-                method=scaling_dict["channel"]["method"],
-                dims=tuple(scaling_dict["channel"]["dims"]),
-            ),
+            target=deserialize_variable_scaling(scaling_dict["target"]),
+            channel=deserialize_variable_scaling(scaling_dict["channel"]),
         )
 
     @classmethod