"""Model serialization utilities for BayesianMMM.
This module provides functions to save and load BayesianMMM models,
separating the serialization logic from the main model class.
"""
from __future__ import annotations
import gzip
import json
import os
import shutil
import tempfile
import warnings
from pathlib import Path
from typing import TYPE_CHECKING, Any
import arviz as az
import numpy as np
from loguru import logger
if TYPE_CHECKING:
from .data_loader import PanelDataset
from .model import BayesianMMM
[docs]
class MMMSerializer:
"""Handles save/load operations for BayesianMMM models.
This class encapsulates all serialization logic, keeping the main
BayesianMMM class focused on modeling.
Examples
--------
>>> # Save a model
>>> MMMSerializer.save(model, "models/my_mmm")
>>> # Load a model
>>> model = MMMSerializer.load("models/my_mmm", panel)
"""
# Version for saved model format. "1.1" adds `garden_ref` /
# `model_class_qualname` (Model Garden: a saved model may be a bespoke
# BayesianMMM subclass, so `load()` reconstructs THAT class, not the base).
_FORMAT_VERSION = "1.1"
@staticmethod
def _is_extended(model: Any) -> bool:
"""True for the BaseExtendedMMM family (Nested / Multivariate /
Combined / StructuralNested), which serializes via its own flavor:
the core save/load contract is panel-shaped (media re-standardization,
panel-compatibility gates) and does not apply to array-level models."""
try:
from .mmm_extensions.models.base import BaseExtendedMMM
return isinstance(model, BaseExtendedMMM)
except Exception: # noqa: BLE001 — extensions not importable
return False
[docs]
@classmethod
def save(
cls,
model: BayesianMMM,
path: str | Path,
save_trace: bool = True,
compress: bool = True,
) -> None:
"""
Save a BayesianMMM model to disk.
Parameters
----------
model : BayesianMMM
The model instance to save.
path : str or Path
Directory path where the model will be saved.
save_trace : bool, default True
Whether to save the fitted trace. Set to False for a smaller
save file if you only need configurations.
compress : bool, default True
Whether to compress the trace file with gzip.
Raises
------
ValueError
If the model has no trace and save_trace is True.
"""
if cls._is_extended(model):
return cls._save_extended(model, path, save_trace, compress)
final = Path(path)
final.parent.mkdir(parents=True, exist_ok=True)
# Write everything into a temp sibling dir, then atomically swap it in.
# A crash mid-save leaves the PREVIOUS model intact rather than a
# half-written directory — important now that models are shared
# cross-project through the Model Garden registry.
work = final.parent / f".{final.name}.tmp"
if work.exists():
shutil.rmtree(work, ignore_errors=True)
work.mkdir(parents=True)
try:
# 1. Save metadata
metadata = cls._collect_metadata(model)
with open(work / "metadata.json", "w") as f:
json.dump(metadata, f, indent=2)
# 2. Save configurations
configs = cls._collect_configs(model)
with open(work / "configs.json", "w") as f:
json.dump(configs, f, indent=2, default=str)
# 3. Save scaling parameters
scaling_params = cls._collect_scaling_params(model)
with open(work / "scaling_params.json", "w") as f:
json.dump(scaling_params, f, indent=2)
# 4. Save trace (if fitted and requested)
if save_trace and model._trace is not None:
cls._save_trace(model._trace, work, compress)
# 5. Save trend features if they exist
cls._save_trend_features(model, work)
# 6. Save seasonality features
cls._save_seasonality_features(model, work)
cls._atomic_swap(work, final)
finally:
if work.exists():
shutil.rmtree(work, ignore_errors=True)
logger.info(f"Model saved to {final}")
@classmethod
def _save_extended(
cls,
model: Any,
path: str | Path,
save_trace: bool = True,
compress: bool = True,
) -> None:
"""Save a BaseExtendedMMM-family model (extended flavor).
The model instance itself (arrays, extension config dataclasses,
masks, y stats — everything except the PyMC graph and the trace,
which ``__getstate__`` strips) rides ``model.pkl`` via cloudpickle,
so loading needs NO panel. The platform-legible sidecars
(``metadata.json`` / ``configs.json``) are written in the same shape
the load/list surfaces already read.
"""
import cloudpickle
final = Path(path)
final.parent.mkdir(parents=True, exist_ok=True)
work = final.parent / f".{final.name}.tmp"
if work.exists():
shutil.rmtree(work, ignore_errors=True)
work.mkdir(parents=True)
try:
metadata = cls._collect_extended_metadata(model)
with open(work / "metadata.json", "w") as f:
json.dump(metadata, f, indent=2, default=str)
configs = cls._collect_extended_configs(model)
with open(work / "configs.json", "w") as f:
json.dump(configs, f, indent=2, default=str)
with open(work / "model.pkl", "wb") as f:
cloudpickle.dump(model, f)
if save_trace and model._trace is not None:
cls._save_trace(model._trace, work, compress)
cls._atomic_swap(work, final)
finally:
if work.exists():
shutil.rmtree(work, ignore_errors=True)
logger.info(f"Extended model saved to {final}")
@classmethod
def _collect_extended_metadata(cls, model: Any) -> dict[str, Any]:
"""Sidecar metadata for an extended-flavor save (self-describing +
readable by saved_model_settings / list_saved_models)."""
mcls = type(model)
# Computed like the core flavor (garden contract), not hardcoded: a
# non-MMM extended subclass must not be mislabeled "mmm" — downstream
# load gates and saved_model_settings key off this field.
from .garden.contract import model_kind as _model_kind
metadata: dict[str, Any] = {
"version": getattr(model, "_VERSION", "1.0"),
"format_version": cls._FORMAT_VERSION,
"model_flavor": "extended",
"model_class_qualname": f"{mcls.__module__}.{mcls.__qualname__}",
"model_kind": _model_kind(model),
"n_obs": int(getattr(model, "n_obs", 0)),
"n_channels": int(getattr(model, "n_channels", 0)),
"channel_names": list(getattr(model, "channel_names", [])),
}
for attr in ("mediator_names", "outcome_names", "control_names"):
val = getattr(model, attr, None)
if val:
metadata[attr] = list(val)
# Fit provenance from the last fit's diagnostics (stamped by
# BaseExtendedMMM.fit) — the artifact stays self-describing about
# whether its uncertainty is calibrated.
diag = getattr(model, "_fit_diagnostics", None) or {}
if diag.get("fit_method") is not None:
metadata["fit_method"] = str(diag["fit_method"])
metadata["approximate"] = bool(diag.get("approximate", False))
if getattr(model, "experiments", None):
try:
metadata["experiments"] = [e.to_dict() for e in model.experiments]
except Exception: # noqa: BLE001 — provenance is best-effort
pass
return metadata
@classmethod
def _collect_extended_configs(cls, model: Any) -> dict[str, Any]:
"""Best-effort JSON view of an extended model's configuration (the
pickle is ground truth; this keeps configs.json readable)."""
import dataclasses
out: dict[str, Any] = {}
ext_cfg = getattr(model, "config", None)
if ext_cfg is not None and dataclasses.is_dataclass(ext_cfg):
try:
out["extension_config"] = dataclasses.asdict(ext_cfg)
except Exception: # noqa: BLE001
out["extension_config"] = repr(ext_cfg)
# Omit (never a string repr) on failure: saved_model_settings reads
# these with dict access, and `<str> or {}` would pass the truthiness
# guard then crash on .get().
mc = getattr(model, "model_config", None)
if mc is not None:
try:
out["model_config"] = mc.model_dump()
except Exception: # noqa: BLE001
pass
tc = getattr(model, "trend_config", None)
if tc is not None:
try:
out["trend_config"] = tc.to_dict()
except Exception: # noqa: BLE001
pass
return out
@classmethod
def _load_extended(
cls,
path: Path,
metadata: dict[str, Any],
rebuild_model: bool,
) -> Any:
"""Load an extended-flavor save: unpickle the instance, reattach the
trace, optionally rebuild the PyMC graph (parity with the core load —
fails fast on environment drift instead of at first use)."""
import cloudpickle
from .mmm_extensions.models.base import BaseExtendedMMM
# Only warn when a version WAS recorded and differs (legacy
# BaseExtendedMMM.save dirs have no metadata at all).
saved_version = metadata.get("version")
if saved_version is not None and saved_version != BaseExtendedMMM._VERSION:
warnings.warn(
f"Extended model was saved with version {saved_version}, "
f"current version is {BaseExtendedMMM._VERSION}. "
"There may be compatibility issues."
)
pkl = path / "model.pkl"
if not pkl.exists():
raise FileNotFoundError(
f"Extended model save at {path} is missing model.pkl"
)
with open(pkl, "rb") as f:
instance = cloudpickle.load(f)
trace = cls._load_trace(path)
if trace is not None:
instance._trace = trace
if rebuild_model:
instance._model = instance._build_model()
logger.info(f"Extended model loaded from {path}")
return instance
@staticmethod
def _atomic_swap(work: Path, final: Path) -> None:
"""Promote the fully-written temp dir ``work`` to ``final``, preserving
the prior model on failure (rename-based, atomic within one filesystem)."""
backup = final.parent / f".{final.name}.bak"
if backup.exists():
shutil.rmtree(backup, ignore_errors=True)
had_prior = final.exists()
if had_prior:
os.replace(final, backup) # move the old model aside
try:
os.replace(work, final) # promote the new model
except Exception:
if had_prior and backup.exists():
os.replace(backup, final) # roll back
raise
finally:
if backup.exists():
shutil.rmtree(backup, ignore_errors=True)
[docs]
@classmethod
def load(
cls,
path: str | Path,
panel: PanelDataset | None = None,
rebuild_model: bool = True,
) -> BayesianMMM:
"""
Load a saved model from disk.
Parameters
----------
path : str or Path
Directory path where the model was saved.
panel : PanelDataset | None
Panel data to use with the loaded model. Must be compatible
with the original data (same channels, controls, dimensions).
REQUIRED for core (BayesianMMM-family) saves; ignored for
extended-flavor saves (BaseExtendedMMM models carry their arrays
in the pickle and need no panel).
rebuild_model : bool, default True
Whether to rebuild the PyMC model. Set to False if you only
need access to the trace and don't need to make predictions.
Returns
-------
BayesianMMM
Loaded model instance with fitted trace (if available).
Raises
------
ValueError
If the panel data is incompatible with the saved model.
FileNotFoundError
If the model files are not found.
"""
# Import here to avoid circular imports
from .model import BayesianMMM, TrendConfig, geometric_adstock_2d
path = Path(path)
if not path.exists():
raise FileNotFoundError(f"Model directory not found: {path}")
# Old-flavor BaseExtendedMMM.save directory (model.pkl, no sidecars):
# delegate to the extended loader with synthesized metadata so legacy
# extension saves just work instead of a bare metadata.json error.
if not (path / "metadata.json").exists() and (path / "model.pkl").exists():
return cls._load_extended(path, {}, rebuild_model)
# 1. Load metadata
with open(path / "metadata.json", "r") as f:
metadata = json.load(f)
# Format gate (both flavors stamp it). Same-major formats load —
# minor bumps are additive by contract; a HIGHER major means the
# artifact was written by an incompatible future serializer and the
# failure should be loud up front, not a KeyError deep in rebuild.
fmt = str(metadata.get("format_version", "1.0"))
try:
fmt_major = int(fmt.split(".")[0])
except ValueError:
fmt_major = 1
our_major = int(cls._FORMAT_VERSION.split(".")[0])
if fmt_major > our_major:
raise ValueError(
f"Model directory was saved with serializer format {fmt}, but "
f"this mmm-framework understands format major {our_major} "
f"(_FORMAT_VERSION={cls._FORMAT_VERSION}). Upgrade "
"mmm-framework to load it."
)
# Extended flavor (BaseExtendedMMM family): panel-free pickle load.
if metadata.get("model_flavor") == "extended":
return cls._load_extended(path, metadata, rebuild_model)
if panel is None:
raise TypeError(
"MMMSerializer.load requires a `panel` for core "
"(BayesianMMM-family) saves — the model is rebuilt against a "
"compatible panel. Only extended-flavor saves load panel-free."
)
# Version check
cls._check_version(metadata)
# 2. Load configurations
with open(path / "configs.json", "r") as f:
configs = json.load(f)
# Reconstruct configs
from .config import ModelConfig
model_config = ModelConfig(**configs["model_config"])
trend_config = TrendConfig.from_dict(configs["trend_config"])
# 3. Validate panel compatibility
cls._validate_panel_compatibility(panel, metadata)
# 4. Create instance
adstock_alphas = metadata.get("adstock_alphas", [0.0, 0.3, 0.5, 0.7, 0.9])
experiments = None
if metadata.get("experiments"):
from .calibration.likelihood import ExperimentMeasurement
experiments = [
ExperimentMeasurement.from_dict(e) for e in metadata["experiments"]
]
# Model Garden: reconstruct the bespoke subclass recorded at save time
# (falls back to BayesianMMM if its source can't be resolved).
model_cls = cls._resolve_model_class(metadata, BayesianMMM)
# Bespoke per-model params: the constructor re-validates the saved dict
# against the resolved class's CONFIG_SCHEMA (defaults/validators applied),
# so a schema added/relaxed since save still loads. None for base models.
instance = model_cls(
panel=panel,
model_config=model_config,
trend_config=trend_config,
adstock_alphas=adstock_alphas,
experiments=experiments,
model_params=metadata.get("model_params"),
)
if metadata.get("garden_ref"):
try:
instance._garden_ref = dict(metadata["garden_ref"])
except Exception: # noqa: BLE001
pass
# Declarative estimands (mirrors the experiments round-trip above).
if metadata.get("declared_estimands"):
from .estimands.spec import Estimand
instance.declared_estimands = [
Estimand.from_dict(e) for e in metadata["declared_estimands"]
]
# Dataset role-mapping schema: re-validate the saved role contract against
# the resolved class's DATASET_SCHEMA (drift-visible, mirroring the
# model_params re-validation) and re-tag the loaded dataset when the saved
# column names match. Best-effort: the panel-derived schema is the
# fallback, so a missing/incompatible saved schema never blocks a load.
saved_schema = metadata.get("dataset_schema")
if saved_schema:
try:
from .config.dataset import DatasetSchema
schema_cls = getattr(model_cls, "DATASET_SCHEMA", None) or DatasetSchema
validated = schema_cls.model_validate(saved_schema)
instance.dataset = instance.dataset.retag(validated)
except Exception: # noqa: BLE001
pass
# 5. Load scaling parameters
with open(path / "scaling_params.json", "r") as f:
scaling_params = json.load(f)
cls._restore_scaling_params(instance, scaling_params)
# The y / media / control re-standardization below is MMM-specific (single
# standardized KPI + per-channel adstock pre-compute). A non-MMM family
# (e.g. a CFA) re-derived its own data in ``_prepare_data`` during
# reconstruction, so skip it.
if metadata.get("model_kind", "mmm") == "mmm":
# Re-standardize y with loaded params. The multiplicative (semi-log)
# model standardizes log(y); y_raw stays original, so re-log here.
y_src = (
np.log(instance.y_raw)
if getattr(instance, "_multiplicative", False)
else instance.y_raw
)
instance.y = (y_src - instance.y_mean) / instance.y_std
# Re-normalize media with loaded max values
for alpha in instance.adstock_alphas:
adstocked = geometric_adstock_2d(instance.X_media_raw, alpha)
normalized = np.zeros_like(adstocked)
for c, ch_name in enumerate(instance.channel_names):
normalized[:, c] = adstocked[:, c] / (
instance._media_max[ch_name] + 1e-8
)
instance.X_media_adstocked[alpha] = normalized
# Re-standardize controls with loaded params
if instance.X_controls_raw is not None and "control_mean" in scaling_params:
instance.X_controls = (
instance.X_controls_raw - instance.control_mean
) / instance.control_std
# 6. Load trend features if present
cls._load_trend_features(instance, path)
# 7. Load seasonality features if present
cls._load_seasonality_features(instance, path)
# 8. Load trace (if available)
trace = cls._load_trace(path)
if trace is not None:
instance._trace = trace
# 9. Build model if requested
if rebuild_model:
instance._model = instance._build_model()
logger.info(f"Model loaded from {path}")
if instance._trace is not None:
logger.info(
f" Trace loaded: {instance._trace.posterior.dims['chain']} chains, "
f"{instance._trace.posterior.dims['draw']} draws"
)
return instance
[docs]
@classmethod
def save_trace_only(
cls,
trace: az.InferenceData,
path: str | Path,
) -> None:
"""
Save only a trace to a file.
Parameters
----------
trace : az.InferenceData
The trace to save.
path : str or Path
File path for the trace (should end in .nc or .nc.gz).
"""
path = Path(path)
if str(path).endswith(".gz"):
# Save compressed
base_path = Path(str(path)[:-3]) # Remove .gz
trace.to_netcdf(str(base_path))
with open(base_path, "rb") as f_in:
with gzip.open(path, "wb") as f_out:
shutil.copyfileobj(f_in, f_out)
base_path.unlink()
else:
trace.to_netcdf(str(path))
logger.info(f"Trace saved to {path}")
[docs]
@classmethod
def load_trace_only(cls, path: str | Path) -> az.InferenceData:
"""
Load a trace from a file.
Parameters
----------
path : str or Path
File path to the trace (.nc or .nc.gz).
Returns
-------
az.InferenceData
The loaded trace.
"""
path = Path(path)
if str(path).endswith(".gz"):
with tempfile.NamedTemporaryFile(suffix=".nc", delete=False) as tmp:
tmp_path = tmp.name
with gzip.open(path, "rb") as f_in:
with open(tmp_path, "wb") as f_out:
shutil.copyfileobj(f_in, f_out)
trace = az.from_netcdf(tmp_path)
os.unlink(tmp_path)
else:
trace = az.from_netcdf(str(path))
logger.info(f"Trace loaded from {path}")
return trace
# =========================================================================
# Private helper methods
# =========================================================================
@classmethod
def _collect_metadata(cls, model: BayesianMMM) -> dict[str, Any]:
"""Collect model metadata for saving."""
# The panel's own channel columns, which may differ from ``channel_names``
# when a model re-expresses the media axis (e.g. the breakout-weighting
# model groups sub-stream columns into virtual channels). Used for the
# panel-compatibility gate so the check is against what the panel actually
# carries, not the model's derived axis. Degrades to ``channel_names``
# when the panel isn't available (they coincide for a plain MMM).
_panel = getattr(model, "panel", None)
panel_channels = (
list(_panel.coords.channels)
if _panel is not None
else list(model.channel_names)
)
# The panel's own control columns, for the same reason as
# ``panel_channels`` above: several features CONSUME a control column
# and strip it from ``model.control_names`` — a reach/frequency
# frequency_column, and price/promo levers. Recording the post-strip
# list and comparing it to the panel's pre-strip set made every such
# model save fine and fail to load. Scoped by the strip mechanism, not
# by any one feature, so it covers both.
panel_controls = (
list(_panel.coords.controls or [])
if _panel is not None
else list(model.control_names)
)
metadata = {
"version": model._VERSION,
"format_version": cls._FORMAT_VERSION,
"n_obs": model.n_obs,
"n_channels": model.n_channels,
"n_controls": model.n_controls,
"n_time_periods": model.n_periods,
"channel_names": model.channel_names,
"panel_channels": panel_channels,
"control_names": model.control_names,
"panel_controls": panel_controls,
"has_geo": model.has_geo,
"has_product": model.has_product,
"adstock_alphas": model.adstock_alphas,
}
if model.has_geo:
metadata["geo_names"] = model.geo_names
if model.has_product:
metadata["product_names"] = model.product_names
# Model Garden provenance: when the model is a bespoke BayesianMMM
# subclass, record the registry ref + qualified class name so load()
# rebuilds the SAME class (and a cold kernel can find its source).
garden_ref = getattr(model, "_garden_ref", None)
if garden_ref:
metadata["garden_ref"] = dict(garden_ref)
cls = type(model)
if cls.__name__ != "BayesianMMM":
metadata["model_class_qualname"] = f"{cls.__module__}.{cls.__qualname__}"
# Garden model family kind — non-MMM families (e.g. a CFA) skip the
# channel/control panel-compatibility match on reload.
from .garden.contract import model_kind as _model_kind
metadata["model_kind"] = _model_kind(model)
# HOW the model was fit — recorded so the artifact is self-describing and
# a reloaded model's saved-settings digest / reports know whether its
# uncertainty is calibrated (NUTS) or approximate (map/advi/pathfinder).
# `fit()` writes the method used back onto model_config, so this is the
# method of the LAST fit (default "nuts" for an unfitted model).
# WHICH PARADIGM produced it. Stamped before `fit_method`, because a
# frequentist fit legitimately has `fit_method=None` and the block below
# would then stamp nothing — leaving a reloaded ridge fit to read as
# Bayesian by omission (absence of `inference_family` means Bayesian).
try:
im = getattr(model.model_config, "inference_method", None)
if im is not None:
metadata["inference_method"] = str(getattr(im, "value", im))
if getattr(model.model_config, "is_frequentist", False):
from .diagnostics import provenance as _prov
metadata["inference_family"] = _prov.FREQUENTIST
metadata["approximate"] = False
diag = getattr(model, "_fit_diagnostics", None) or {}
for key in (
"estimator",
"interval_kind",
"interval_semantics",
"selection_criterion",
"n_boot",
"block_length",
"effective_dof",
"penalty",
):
if key in diag:
metadata[key] = diag[key]
except Exception: # noqa: BLE001
pass
try:
fm = getattr(model.model_config, "fit_method", None)
fit_method = getattr(fm, "value", fm)
if fit_method is not None:
metadata["fit_method"] = str(fit_method)
# NUTS and SMC are exact (FitMethod.is_approximate); fall back
# to the legacy != "nuts" rule only for unknown method strings.
try:
from .config.enums import FitMethod as _FM
metadata["approximate"] = _FM(
str(fit_method).lower()
).is_approximate
except ValueError:
metadata["approximate"] = str(fit_method).lower() != "nuts"
except Exception: # noqa: BLE001
pass
# Experiment calibration likelihoods (so a reloaded model can be re-fit
# with the same incrementality anchoring it was originally built with).
if getattr(model, "experiments", None):
metadata["experiments"] = [e.to_dict() for e in model.experiments]
# Declarative estimands associated with the model (the counterfactual
# causal lens); round-tripped with a schema_version so drift is visible.
if getattr(model, "declared_estimands", None):
metadata["declared_estimands"] = [
e.to_dict() for e in model.declared_estimands
]
# Bespoke per-model parameters (validated against the model's
# CONFIG_SCHEMA). Stored as a plain dict so the reloaded model re-validates
# it against the (possibly evolved) schema. ``schema_version`` makes drift
# visible; the model class owns its CONFIG_SCHEMA shape.
mp = getattr(model, "model_params", None)
if mp is not None:
from pydantic import BaseModel as _BaseModel
metadata["model_params"] = (
mp.model_dump() if isinstance(mp, _BaseModel) else dict(mp)
)
schema = getattr(type(model), "CONFIG_SCHEMA", None)
metadata["model_params_schema_version"] = (
getattr(schema, "SCHEMA_VERSION", 1) if schema is not None else None
)
# Dataset role-mapping schema (the flexible data layer): record the role
# contract the model actually used so a reloaded — or garden-listed —
# model carries its data needs, re-validated on load against the class's
# DATASET_SCHEMA. Present for every model (auto-derived for plain MMM).
ds = getattr(model, "dataset", None)
ds_schema = getattr(ds, "schema", None)
if ds_schema is not None:
try:
metadata["dataset_schema"] = ds_schema.model_dump(mode="json")
dsc = getattr(type(model), "DATASET_SCHEMA", None)
metadata["dataset_schema_version"] = getattr(
ds_schema,
"SCHEMA_VERSION",
getattr(dsc, "SCHEMA_VERSION", "1.0"),
)
except Exception: # noqa: BLE001
pass
return metadata
@classmethod
def _collect_configs(cls, model: BayesianMMM) -> dict[str, Any]:
"""Collect model configurations for saving."""
return {
"model_config": model.model_config.model_dump(),
"trend_config": model.trend_config.to_dict(),
"mff_config": model.mff_config.model_dump(),
}
@classmethod
def _collect_scaling_params(cls, model: BayesianMMM) -> dict[str, Any]:
"""Collect scaling parameters for saving."""
scaling_params = {
"y_mean": model.y_mean,
"y_std": model.y_std,
"media_max": {k: float(v) for k, v in model._media_max.items()},
}
if model.X_controls_raw is not None:
scaling_params["control_mean"] = model.control_mean.tolist()
scaling_params["control_std"] = model.control_std.tolist()
return scaling_params
@classmethod
def _save_trace(
cls,
trace: az.InferenceData,
path: Path,
compress: bool,
) -> None:
"""Save trace to disk, optionally compressed."""
trace_path = path / "trace.nc"
trace.to_netcdf(str(trace_path))
if compress:
with open(trace_path, "rb") as f_in:
with gzip.open(str(trace_path) + ".gz", "wb") as f_out:
shutil.copyfileobj(f_in, f_out)
trace_path.unlink() # Remove uncompressed file
@classmethod
def _save_trend_features(cls, model: BayesianMMM, path: Path) -> None:
"""Save trend features if they exist."""
trend_features_to_save = {}
for key, value in model.trend_features.items():
if isinstance(value, np.ndarray):
trend_features_to_save[key] = value.tolist()
elif isinstance(value, dict):
trend_features_to_save[key] = value
else:
trend_features_to_save[key] = value
if trend_features_to_save:
with open(path / "trend_features.json", "w") as f:
json.dump(trend_features_to_save, f, indent=2)
@classmethod
def _save_seasonality_features(cls, model: BayesianMMM, path: Path) -> None:
"""Save seasonality features."""
season_features_to_save = {}
for key, value in model.seasonality_features.items():
if isinstance(value, np.ndarray):
season_features_to_save[key] = value.tolist()
if season_features_to_save:
with open(path / "seasonality_features.json", "w") as f:
json.dump(season_features_to_save, f, indent=2)
@classmethod
def _resolve_model_class(cls, metadata: dict[str, Any], default: type) -> type:
"""Resolve the class to reconstruct on load.
Model Garden models are bespoke ``BayesianMMM`` subclasses; their source
is imported from the ``garden_ref`` recorded at save time. If the source
can't be found (e.g. loaded in a different session), fall back to the
base class with a warning — the trace + scaling still load for read-only
inspection, the custom build hooks are just not reapplied.
"""
ref = metadata.get("garden_ref")
if ref:
# Garden-registered models resolve via their stored source; if that
# can't be found (e.g. a different session), degrade to ``default``.
try:
from .garden.loader import load_garden_class_from_path
return load_garden_class_from_path(
ref.get("source_path"), ref.get("class_name")
)
except Exception as exc: # noqa: BLE001
warnings.warn(
f"Could not load garden model class {ref.get('name')!r} "
f"({exc}); falling back to {default.__name__} for trace "
"inspection."
)
return default
# No garden_ref: a bespoke ``BayesianMMM`` subclass constructed directly
# (e.g. a CFA). Import it by its recorded fully-qualified name so it
# round-trips whenever its module is importable; else degrade.
qualname = metadata.get("model_class_qualname")
if qualname and "." in qualname:
module_path, _, cls_name = qualname.rpartition(".")
try:
import importlib
candidate = getattr(
importlib.import_module(module_path), cls_name, None
)
if isinstance(candidate, type):
return candidate
except Exception as exc: # noqa: BLE001
warnings.warn(
f"Could not import model class {qualname!r} ({exc}); "
f"falling back to {default.__name__} for trace inspection."
)
return default
@classmethod
def _check_version(cls, metadata: dict[str, Any]) -> None:
"""Check version compatibility and warn if needed."""
from .model import BayesianMMM
saved_version = metadata.get("version", "0.0.0")
if saved_version != BayesianMMM._VERSION:
warnings.warn(
f"Model was saved with version {saved_version}, "
f"current version is {BayesianMMM._VERSION}. "
"There may be compatibility issues."
)
@classmethod
def _validate_panel_compatibility(
cls,
panel: PanelDataset,
metadata: dict[str, Any],
) -> None:
"""Validate that the panel is compatible with the saved model."""
# Channel/control identity is MMM-specific. A non-MMM family (e.g. a CFA)
# carries an indicator matrix, not channels — it sets ``channel_names=[]``
# while the panel still lists its observed columns, so skip the match.
if metadata.get("model_kind", "mmm") != "mmm":
return
# Compare against the panel's own channels (``panel_channels``) when the
# save recorded them — a model may re-express its media axis (e.g. the
# breakout-weighting model's virtual channels), so ``channel_names`` need
# not equal the panel columns. Old saves (no ``panel_channels``) fall back
# to ``channel_names``, which equals the panel columns for a plain MMM.
expected = metadata.get("panel_channels", metadata["channel_names"])
if panel.coords.channels != expected:
raise ValueError(
f"Panel channels {panel.coords.channels} don't match "
f"saved model channels {expected}"
)
# Compare against the panel's own control columns, not the model's
# post-strip axis. Falls back to ``control_names`` for saves written
# before ``panel_controls`` existed, where the two coincide unless a
# column was consumed — those old saves stay exactly as loadable (or
# not) as they were.
expected_controls = metadata.get("panel_controls", metadata["control_names"])
if list(panel.coords.controls or []) != list(expected_controls or []):
raise ValueError(
f"Panel controls {panel.coords.controls} don't match "
f"saved model controls {expected_controls}"
)
@classmethod
def _restore_scaling_params(
cls,
instance: BayesianMMM,
scaling_params: dict[str, Any],
) -> None:
"""Restore scaling parameters to the model instance."""
instance.y_mean = scaling_params["y_mean"]
instance.y_std = scaling_params["y_std"]
instance._media_max = scaling_params["media_max"]
instance._scaling_params["y_mean"] = scaling_params["y_mean"]
instance._scaling_params["y_std"] = scaling_params["y_std"]
instance._scaling_params["media_max"] = scaling_params["media_max"]
if "control_mean" in scaling_params:
instance.control_mean = np.array(scaling_params["control_mean"])
instance.control_std = np.array(scaling_params["control_std"])
instance._scaling_params["control_mean"] = instance.control_mean
instance._scaling_params["control_std"] = instance.control_std
@classmethod
def _load_trend_features(cls, instance: BayesianMMM, path: Path) -> None:
"""Load trend features if present."""
trend_features_path = path / "trend_features.json"
if trend_features_path.exists():
with open(trend_features_path, "r") as f:
trend_features = json.load(f)
# Convert lists back to numpy arrays
for key, value in trend_features.items():
if isinstance(value, list):
instance.trend_features[key] = np.array(value)
else:
instance.trend_features[key] = value
@classmethod
def _load_seasonality_features(cls, instance: BayesianMMM, path: Path) -> None:
"""Load seasonality features if present."""
season_features_path = path / "seasonality_features.json"
if season_features_path.exists():
with open(season_features_path, "r") as f:
season_features = json.load(f)
for key, value in season_features.items():
if isinstance(value, list):
instance.seasonality_features[key] = np.array(value)
@classmethod
def _load_trace(cls, path: Path) -> az.InferenceData | None:
"""Load trace if available."""
trace_path_gz = path / "trace.nc.gz"
trace_path = path / "trace.nc"
if trace_path_gz.exists():
with tempfile.NamedTemporaryFile(suffix=".nc", delete=False) as tmp:
tmp_path = tmp.name
with gzip.open(trace_path_gz, "rb") as f_in:
with open(tmp_path, "wb") as f_out:
shutil.copyfileobj(f_in, f_out)
trace = az.from_netcdf(tmp_path)
# Materialize the trace into memory BEFORE deleting the temp file —
# az.from_netcdf can be lazy (engine-dependent), and returning a lazy
# InferenceData that references a unlinked temp makes any later access
# raise "Unable to open file" (hit when reloading a model in a fresh
# kernel for interpretation).
try:
for _group in getattr(trace, "_groups_all", []):
getattr(trace, _group).load()
except Exception:
pass
os.unlink(tmp_path) # Clean up temp file
return trace
elif trace_path.exists():
return az.from_netcdf(str(trace_path))
return None