Refactor ModelBuilder into smaller classes (#1870)

* break out into smaller classes

* change the default to True

* WIP unit testing

* requires_model methods

* test_graphviz

* TODOs and fix MOdelBuilder tests

* remove test model file

* fix customer and MMM tests

* WIP fix clv tests

* fix clv base tests

* fix remaining clv tests

* docstrings

* docstring edit

* remove idata_to_init_kwargs

* CLV models inherit from BaseModelBuilder

* WIP test cleanup

* fix base model tests

* clean up tests

* testing coverage

* add abstract methods to ModelBuilder

* docstrings

* docstrings return model_config

* remove duplicated self.data attr

---------

Co-authored-by: Will Dean <[email protected]>
Co-authored-by: Juan Orduz <[email protected]>

Author: 3 people

pymc_marketing/clv/models/basic.py

@@ -13,10 +13,8 @@
 #   limitations under the License.
 """CLV Model base class."""
 
-import json
 import warnings
 from collections.abc import Sequence
-from pathlib import Path
 from typing import Literal, cast
 
 import arviz as az
@@ -27,9 +25,8 @@
 from pymc.backends.base import MultiTrace
 from pymc.model.core import Model
 
-from pymc_marketing.model_builder import ModelBuilder
+from pymc_marketing.model_builder import DifferentModelError, ModelBuilder
 from pymc_marketing.model_config import ModelConfig, parse_model_config
-from pymc_marketing.utils import from_netcdf
 
 
 class CLVModel(ModelBuilder):
@@ -46,6 +43,7 @@ def __init__(
         sampler_config: dict | None = None,
         non_distributions: list[str] | None = None,
     ):
+        self.data = data
         model_config = model_config or {}
 
         deprecated_keys = [key for key in model_config if key.endswith("_prior")]
@@ -60,14 +58,14 @@ def __init__(
 
             model_config[new_key] = model_config.pop(key)
 
-        model_config = parse_model_config(
-            model_config,
+        super().__init__(model_config, sampler_config)
+
+        # Parse model config after merging with defaults
+        self.model_config = parse_model_config(
+            self.model_config,
             non_distributions=non_distributions,
         )
 
-        super().__init__(model_config, sampler_config)
-        self.data = data
-
     @staticmethod
     def _validate_cols(
         data: pd.DataFrame,
@@ -260,59 +258,39 @@ def _fit_approx(
             )
 
     @classmethod
-    def load(cls, fname: str):
-        """Create a ModelBuilder instance from a file.
-
-        Loads inference data for the model.
-
-        Parameters
-        ----------
-        fname : string
-            This denotes the name with path from where idata should be loaded from.
+    def idata_to_init_kwargs(cls, idata: az.InferenceData) -> dict:
+        """Create the initialization kwargs from an InferenceData object."""
+        kwargs = cls.attrs_to_init_kwargs(idata.attrs)
+        kwargs["data"] = idata.fit_data.to_dataframe()
 
-        Returns
-        -------
-        Returns an instance of ModelBuilder.
-
-        Raises
-        ------
-        ValueError
-            If the inference data that is loaded doesn't match with the model.
-
-        Examples
-        --------
-        >>> class MyModel(ModelBuilder):
-        >>>     ...
-        >>> name = "./mymodel.nc"
-        >>> imported_model = MyModel.load(name)
-
-        """
-        filepath = Path(str(fname))
-        idata = from_netcdf(filepath)
-        return cls._build_with_idata(idata)
+        return kwargs
 
     @classmethod
-    def _build_with_idata(cls, idata: az.InferenceData):
-        dataset = idata.fit_data.to_dataframe()
+    def build_from_idata(cls, idata: az.InferenceData) -> None:
+        """Build the model from the InferenceData object."""
+        kwargs = cls.idata_to_init_kwargs(idata)
         with warnings.catch_warnings():
             warnings.filterwarnings(
                 "ignore",
                 category=DeprecationWarning,
             )
-            model = cls(
-                dataset,
-                model_config=json.loads(idata.attrs["model_config"]),  # type: ignore
-                sampler_config=json.loads(idata.attrs["sampler_config"]),
-            )
+            model = cls(**kwargs)
 
         model.idata = idata
         model._rename_posterior_variables()
 
         model.build_model()  # type: ignore
         if model.id != idata.attrs["id"]:
-            raise ValueError(f"Inference data not compatible with {cls._model_type}")
+            msg = (
+                "The model id in the InferenceData does not match the model id. "
+                "There was no error loading the inference data, but the model may "
+                "be different. "
+                "Investigate if the model structure or configuration has changed."
+            )
+            raise DifferentModelError(msg)
         return model
 
+    # TODO: Remove in 2026Q1?
     def _rename_posterior_variables(self):
         """Rename variables in the posterior group to remove the _prior suffix.
 
@@ -355,7 +333,7 @@ def thin_fit_result(self, keep_every: int):
         self.fit_result  # noqa: B018 (Raise Error if fit didn't happen yet)
         assert self.idata is not None  # noqa: S101
         new_idata = self.idata.isel(draw=slice(None, None, keep_every)).copy()
-        return type(self)._build_with_idata(new_idata)
+        return self.build_from_idata(new_idata)
 
     @property
     def default_sampler_config(self) -> dict:
@@ -378,8 +356,3 @@ def fit_summary(self, **kwargs):
             return res["mean"].rename("value")
         else:
             return az.summary(self.fit_result, **kwargs)
-
-    @property
-    def output_var(self):
-        """Output variable of the model."""
-        pass

pymc_marketing/customer_choice/mnl_logit.py

@@ -25,13 +25,13 @@
 import pytensor.tensor as pt
 from pymc_extras.prior import Prior
 
-from pymc_marketing.model_builder import ModelBuilder
+from pymc_marketing.model_builder import RegressionModelBuilder
 from pymc_marketing.model_config import parse_model_config
 
 HDI_ALPHA = 0.5
 
 
-class MNLogit(ModelBuilder):
+class MNLogit(RegressionModelBuilder):
     """
     Multinomial Logit class.
 

pymc_marketing/customer_choice/mv_its.py

@@ -27,13 +27,13 @@
 from xarray import DataArray
 
 from pymc_marketing.mmm.additive_effect import MuEffect
-from pymc_marketing.model_builder import ModelBuilder, create_idata_accessor
+from pymc_marketing.model_builder import RegressionModelBuilder, create_idata_accessor
 from pymc_marketing.model_config import parse_model_config
 
 HDI_ALPHA = 0.5
 
 
-class MVITS(ModelBuilder):
+class MVITS(RegressionModelBuilder):
     """Multivariate Interrupted Time Series class.
 
     Class to perform a multivariate interrupted time series analysis with the
@@ -251,7 +251,7 @@ def _generate_and_preprocess_model_data(
             ],
         }
 
-    def build_model(
+    def build_model(  # type: ignore[override]
         self,
         X: pd.DataFrame,
         y: pd.Series | np.ndarray,

pymc_marketing/customer_choice/nested_logit.py

@@ -26,13 +26,13 @@
 from pymc_extras.prior import Prior
 from pytensor.tensor.variable import TensorVariable
 
-from pymc_marketing.model_builder import ModelBuilder
+from pymc_marketing.model_builder import RegressionModelBuilder
 from pymc_marketing.model_config import parse_model_config
 
 HDI_ALPHA = 0.5
 
 
-class NestedLogit(ModelBuilder):
+class NestedLogit(RegressionModelBuilder):
     """
     Nested Logit class.
 

pymc_marketing/mmm/base.py

@@ -45,14 +45,14 @@
     ValidateDateColumn,
     ValidateTargetColumn,
 )
-from pymc_marketing.model_builder import ModelBuilder
+from pymc_marketing.model_builder import RegressionModelBuilder
 
 __all__ = ["BaseValidateMMM", "MMMModelBuilder"]
 
 from pydantic import Field, validate_call
 
 
-class MMMModelBuilder(ModelBuilder):
+class MMMModelBuilder(RegressionModelBuilder):
     """Base class for Marketing Mix Models (MMM)."""
 
     model: pm.Model

pymc_marketing/mmm/mmm.py

@@ -503,7 +503,7 @@ def forward_pass(self, x: pt.TensorVariable | npt.NDArray) -> pt.TensorVariable:
 
         return second.apply(x=first.apply(x=x, dims="channel"), dims="channel")
 
-    def build_model(
+    def build_model(  # type: ignore[override]
         self,
         X: pd.DataFrame,
         y: pd.Series | np.ndarray,

pymc_marketing/mmm/multidimensional.py

@@ -62,7 +62,10 @@
     add_noise_to_channel_allocation,
     create_zero_dataset,
 )
-from pymc_marketing.model_builder import ModelBuilder, _handle_deprecate_pred_argument
+from pymc_marketing.model_builder import (
+    RegressionModelBuilder,
+    _handle_deprecate_pred_argument,
+)
 from pymc_marketing.model_config import parse_model_config
 from pymc_marketing.model_graph import deterministics_to_flat
 
@@ -75,7 +78,7 @@
 warnings.warn(warning_msg, FutureWarning, stacklevel=1)
 
 
-class MMM(ModelBuilder):
+class MMM(RegressionModelBuilder):
     """Marketing Mix Model class for estimating the impact of marketing channels on a target variable.
 
     This class implements the core functionality of a Marketing Mix Model (MMM), allowing for the