MMM Plot Suite Migration: arviz_plots Multi-Backend Support #2098

isofer · 2025-11-21T21:08:12Z

MMM Plot Suite Migration: arviz_plots Multi-Backend Support

Summary

This PR introduces a new version of MMMPlotSuite built on arviz_plots, enabling multi-backend plotting (matplotlib, plotly, bokeh) while preserving full backward compatibility. The original matplotlib-only implementation is preserved as LegacyMMMPlotSuite and remains the default in v0.18.0, with users able to opt into the new suite via configuration.
This closes #2083 and #2081 (duplicated).

Key highlights:

✨ New arviz_plots-based plotting suite with multi-backend support
🔄 Original matplotlib suite preserved as LegacyMMMPlotSuite
⚙️ Configuration system for version control and backend selection
✅ Comprehensive test coverage (~3000+ lines)
🔧 Zero breaking changes (backward compatible by default)

Key Changes

🎨 New Multi-Backend Plotting Suite (`plot.py`)

Complete rewrite of MMMPlotSuite using arviz_plots:

Multi-backend support: matplotlib, plotly, and bokeh
Returns: PlotCollection objects instead of (Figure, Axes) tuples
All plot methods migrated:
- posterior_predictive() - Posterior predictive time series
- contributions_over_time() - Variable contributions over time with HDI
- saturation_scatterplot() - Channel saturation scatter plots
- saturation_curves() - Saturation curves with samples and HDI overlays
- budget_allocation_roas() - Budget allocation ROAS distributions
- allocated_contribution_by_channel_over_time() - Allocated contributions
- sensitivity_analysis() - Sensitivity analysis results
- uplift_curve() - Uplift curves
- marginal_curve() - Marginal effects

Architectural improvements:

Data parameter standardization: All methods accept explicit data parameters with fallback to self.idata
Backend resolution: _resolve_backend() handles global config with per-method override
Cleaner helpers: Removed matplotlib-specific internal methods like _init_subplots(), _build_subplot_title(), _add_median_and_hdi()

🔄 Backward Compatibility & Legacy Suite

Legacy Suite Created (legacy_plot.py, 1,937 lines):

New file: Complete copy of original matplotlib-only implementation
Class name: LegacyMMMPlotSuite
Functionality: Unchanged - maintains exact matplotlib-only behavior
Status: Deprecated, removal planned for v0.20.0

Version Switching (multidimensional.py):
Modified MMM.plot property to support version control:

@property
def plot(self):
    if mmm_config.get("plot.use_v2", False):
        return MMMPlotSuite(idata=self.idata)  # New arviz_plots suite
    else:
        # Shows deprecation warning
        return LegacyMMMPlotSuite(idata=self.idata)  # Legacy matplotlib suite

Deprecation Timeline:

v0.18.0 (this release): Legacy default, new suite opt-in
v0.19.0 (planned): New suite becomes default
v0.20.0 (planned): Legacy suite removed

⚙️ Configuration System (`config.py` - NEW)

Added global MMM configuration (156 lines):

from pymc_marketing.mmm import mmm_config

# Configuration options:
mmm_config["plot.backend"] = "matplotlib" | "plotly" | "bokeh"  # Default: "matplotlib"
mmm_config["plot.use_v2"] = False  # True = new suite, False = legacy (default)
mmm_config["plot.show_warnings"] = True  # Control deprecation warnings

Features:

Backend validation with helpful warnings
Modeled after ArviZ's rcParams pattern
Per-method override via backend parameter on all plot methods
Exported via pymc_marketing.mmm for easy access

📊 API Improvements

Data parameter standardization: All methods accept explicit data parameters with fallback to self.idata
- contributions_over_time(data=...)
- saturation_scatterplot(constant_data=..., posterior_data=...)
- saturation_curves(constant_data=..., posterior_data=...)
- sensitivity_analysis(data=...), uplift_curve(data=...), marginal_curve(data=...)

🧪 Comprehensive Testing (~3,100 lines of test code)

New Test Files:

tests/mmm/conftest.py (307 lines): Shared fixtures for all plotting tests
tests/mmm/test_plot.py (1,719 lines): Rewritten for new suite
- All methods tested with new arviz_plots implementation
- PlotCollection return type validation
- Backend parameter tests
tests/mmm/test_legacy_plot.py (1,054 lines): Full legacy suite coverage
- Complete test suite for matplotlib-only implementation
- Ensures legacy behavior preserved
tests/mmm/test_plot_compatibility.py (475 lines): Version switching tests
- Tests mmm_config["plot.use_v2"] switching
- Validates correct suite returned
- Deprecation warning tests
tests/mmm/test_plot_data_parameters.py (120 lines): Data parameter validation
tests/mmm/test_legacy_plot_imports.py (33 lines): Import validation
tests/mmm/test_legacy_plot_regression.py (56 lines): Legacy behavior regression tests

Breaking Changes

Return Types

Legacy (matplotlib-only):

fig, axes = mmm.plot.posterior_predictive()

New (multi-backend):

pc = mmm.plot.posterior_predictive()
pc.show()  # Display
pc.save("plot.png")  # Save

Parameter Changes

posterior_predictive(var=...): Now accepts str instead of list[str] (call multiple times for multiple variables)
All methods: Removed ax, figsize, colors, subplot_kwargs, rc_params parameters (use PlotCollection API instead)
All methods: Added backend parameter for per-method override

Method Changes

budget_allocation(): Now a deprecated wrapper that calls budget_allocation_roas() with deprecation warning
budget_allocation_roas(): New primary method for budget allocation visualizations

Migration Guide

Quick Start

from pymc_marketing.mmm import mmm_config

# Enable new plotting suite
mmm_config["plot.use_v2"] = True

# Set global backend
mmm_config["plot.backend"] = "plotly"

# Use plots
pc = mmm.plot.posterior_predictive()
pc.show()

Opting Out (Temporary)

# Use legacy suite (will show deprecation warning)
mmm_config["plot.use_v2"] = False
fig, ax = mmm.plot.posterior_predictive()

Files Changed

Statistics: 15 files changed, 9,092 insertions(+), 1,886 deletions(-)

New Files

pymc_marketing/mmm/config.py (156 lines)
pymc_marketing/mmm/legacy_plot.py (1,937 lines)
tests/mmm/conftest.py (307 lines)
tests/mmm/test_legacy_plot.py (1,054 lines)
tests/mmm/test_legacy_plot_imports.py (33 lines)
tests/mmm/test_legacy_plot_regression.py (56 lines)
tests/mmm/test_plot_compatibility.py (475 lines)
tests/mmm/test_plot_data_parameters.py (120 lines)
CLAUDE.md (272 lines)

Modified Files

pymc_marketing/mmm/__init__.py (+2 lines: export mmm_config)
pymc_marketing/mmm/plot.py (complete rewrite: ~2660 lines changed)
pymc_marketing/mmm/multidimensional.py (+54 lines: version control logic)
tests/mmm/test_plot.py (heavily rewritten: 1719 lines)
.gitignore (+3 lines)
CLAUDE.md (NEW, 272 lines: project documentation)

Testing

✅ All new suite methods tested with all 3 backends (matplotlib, plotly, bokeh)
✅ Backend configuration and override tests
✅ Version switching and deprecation warning tests
✅ Return type compatibility tests
✅ Data parameter functionality tests
✅ Legacy suite tests preserved and passing

Documentation

CLAUDE.md: Added comprehensive MMM plotting architecture documentation
config.py: Extensive docstrings with usage examples
Method docstrings: Updated with:
- PlotCollection return types
- backend parameter documentation
- Version directives (.. versionadded::, .. versionchanged::)
- Data parameter documentation
Deprecation warnings: Include migration timeline and guide references

Timeline

v0.18.0 (this PR): Introduce new suite with legacy as default + deprecation warning
v0.19.0: Switch default to new suite (plot.use_v2 = True)
v0.20.0: Remove legacy suite entirely

Checklist

Notes

Zero breaking changes for existing users (legacy suite remains default)
Opt-in migration via mmm_config["plot.use_v2"] = True
Clear deprecation timeline with helpful warnings
Git history preserved for renamed files (use git log --follow)

📚 Documentation preview 📚: https://pymc-marketing--2098.org.readthedocs.build/en/2098/

williambdean · 2025-11-21T22:28:23Z

CLAUDE.md

remove this

We already added claude commands and skills to the repo so it seems appropriate to add this file.

Can it be renamed to AGENTS.md then if this if general enough then.

Claude.md, like the skills and commands that were already uploaded, are used by Claude code. It expect for a with that name and these set of folders.
The text in Claude.md like the skills and command is very generic and can be used by any AI Platform, but we need this specific structure to get the most of Claude code.
So I don't think it make sense to change the name of this file while leaving all the .claude folder.

The current status of the repo is weird, we uploaded skills and command which mean we support Claude Code, but we are missing Claude.md, so Claude Code cannot run on this repo.
I think we should either go all-in and support claude code and add the Claude.md file, or remove the .claude folder. I prefer to add Claude.md

pymc_marketing/mmm/plot.py

codecov · 2025-11-22T19:36:38Z

Codecov Report

❌ Patch coverage is 77.61905% with 141 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.65%. Comparing base (4983c1e) to head (e9a05ba).

Files with missing lines	Patch %	Lines
pymc_marketing/mmm/legacy_plot.py	76.61%	141 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2098      +/-   ##
==========================================
+ Coverage   92.57%   92.65%   +0.07%     
==========================================
  Files          69       71       +2     
  Lines        9430     9760     +330     
==========================================
+ Hits         8730     9043     +313     
- Misses        700      717      +17

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

cetagostini

Just a few docstrings changes.

cetagostini · 2025-11-24T06:20:41Z

pymc_marketing/mmm/plot.py

+        --------
+        Basic usage:
+
+        >>> mmm.sample_posterior_predictive(X)


We avoid this and use code blocks.

cetagostini · 2025-11-24T06:23:04Z

pymc_marketing/mmm/plot.py

+        constant_data : xr.Dataset, optional
+            Dataset containing constant_data group. If None, uses self.idata.constant_data.
+
+            .. versionadded:: 0.18.0


Duplicated text?

I decided to simply remove all of these .. versionadded::. I don't think we need them here.

pymc_marketing/mmm/plot.py

williambdean · 2025-11-26T06:58:43Z

CLAUDE.md

Can it be renamed to AGENTS.md then if this if general enough then.

williambdean · 2025-11-26T07:06:54Z

tests/mmm/test_legacy_plot_imports.py

+def test_legacy_plot_module_exists():
+    """Test that legacy_plot module exists and can be imported."""
+    try:
+        from pymc_marketing.mmm import legacy_plot
+
+        assert hasattr(legacy_plot, "LegacyMMMPlotSuite")
+    except ImportError as e:
+        pytest.fail(f"Failed to import legacy_plot: {e}")
+
+
+def test_legacy_class_name():
+    """Test that legacy suite class is named LegacyMMMPlotSuite."""
+    from pymc_marketing.mmm.legacy_plot import LegacyMMMPlotSuite
+
+    assert LegacyMMMPlotSuite.__name__ == "LegacyMMMPlotSuite"


remove. These tests will be obvious failures based on other tests.

williambdean · 2025-11-26T07:11:39Z

pymc_marketing/mmm/plot.py

+    def _resolve_backend(self, backend: str | None) -> str:
+        """Resolve backend parameter to actual backend string."""
+        from pymc_marketing.mmm.config import mmm_config
+
+        return backend or mmm_config["plot.backend"]


Why not toplevel import?

sorry, agent generated code. I didn't catch it. fixed

pymc_marketing/mmm/multidimensional.py

williambdean · 2025-11-26T07:16:38Z

tests/mmm/test_plot_data_parameters.py

+def test_contributions_over_time_accepts_data_parameter(mock_posterior_data):
+    """Test that contributions_over_time accepts data parameter."""
+    # Create suite without idata
+    suite = MMMPlotSuite(idata=None)
+
+    # Should work with explicit data parameter
+    pc = suite.contributions_over_time(var=["intercept"], data=mock_posterior_data)
+
+    assert isinstance(pc, arviz_plots.PlotCollection)
+
+
+def test_contributions_over_time_data_parameter_fallback(mock_idata_with_posterior):
+    """Test that contributions_over_time falls back to self.idata.posterior."""
+    suite = MMMPlotSuite(idata=mock_idata_with_posterior)
+
+    # Should work without data parameter (fallback)
+    pc = suite.contributions_over_time(var=["intercept"])
+


@copilot use pytest.mark.parametrize to combine these two tests

I think these are fine to keep separated.

Copilot · 2025-11-26T07:17:51Z

@williambdean I've opened a new pull request, #2103, to work on those changes. Once the pull request is ready, I'll request review from you.

williambdean · 2025-11-26T07:20:35Z

pymc_marketing/mmm/plot.py

+        raise NotImplementedError(
+            "budget_allocation() was removed in MMMPlotSuite v2.\n\n"
+            "The new arviz_plots-based implementation doesn't support this chart type.\n\n"
+            "Alternatives:\n"
+            "  1. For ROI distributions: use budget_allocation_roas()\n"
+            "  2. To use old method: set mmm_config['plot.use_v2'] = False\n"
+            "  3. Implement custom bar chart using the samples data\n\n"
+            "See documentation: https://docs.pymc-marketing.io/en/latest/mmm/plotting_migration.html#budget-allocation"
+        )


Is the goal to remove these missing methods entirely in 0.20.0 ? Or will they only exist for the matplotlib backend ?

CC: @cetagostini

I've consulted with @cetagostini regarding this change, and followed his recommendation.
His point was that the current implementation of this graph is not the best, and so we decided to create budget_allocation_roas and deprecated budget_allocation entirely.
in 0.20.0 we will remove LegacyMMMPlotSuite and only the new suite will be available.

isofer · 2025-11-27T22:57:54Z

@williambdean @cetagostini
I've replied/ addressed all comments.
One change that I just pushed is renaming the MMMConfig to MMMPlotConfig and also the instance of mmm_config to mmm_plot_config. The reason was simple - this is a configuration dictionary that affects only plotting but the name implied that it much more general. I don't want people think it affects the model itself or other things that it doesn't affect.

* Initial plan * Combine contributions_over_time tests using pytest.mark.parametrize Co-authored-by: williambdean <57733339+williambdean@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: williambdean <57733339+williambdean@users.noreply.github.com> Co-authored-by: Juan Orduz <juanitorduz@gmail.com>

isofer added 18 commits November 12, 2025 16:00

completed phase 1

b2b881e

adds plan

6dbd6ce

add research

a6ec594

add arviz to 3 functions

3fe9969

a lot of changes to plot.py

278e9bb

clean up for plot.py

d6331a0

doing research

f5574d3

change to research

e140362

finished milestone 1.4

ef02276

finished 3.2

2632822

working on 3.3

dc67ab0

done milestone 4

84eadc9

add check for valid keys in mmm_config

ddce5b9

small change to commit

f354005

small update

26da1d4

Merge branch 'main' into feature/mmmplotsuite-arviz

b36d0e9

adds test_legacy_plot

d5d4bae

remove old test

c321039

github-actions bot added MMM tests labels Nov 21, 2025

small changes

2d4723f

isofer marked this pull request as ready for review November 21, 2025 21:14

isofer requested review from OriolAbril and cetagostini November 21, 2025 21:14

williambdean reviewed Nov 21, 2025

View reviewed changes

pymc_marketing/mmm/plot.py Show resolved Hide resolved

update Claude.md

cdb4877

github-actions bot added the plots label Nov 22, 2025

isofer added 2 commits November 22, 2025 14:25

update env

3741bb2

update pyproject.toml

8cb9d07

update pyproject.toml

cb9dccc

isofer and others added 3 commits November 22, 2025 14:41

Merge branch 'main' into feature/mmmplotsuite-arviz

b612ac8

increase coverage of config.py to 100

9728128

Merge branch 'main' into feature/mmmplotsuite-arviz

599d1f4

cetagostini reviewed Nov 24, 2025

View reviewed changes

williambdean reviewed Nov 26, 2025

View reviewed changes

Copilot AI mentioned this pull request Nov 26, 2025

Address feedback: Combine tests using pytest.mark.parametrize #2103

Merged

6 tasks

Merge branch 'main' into feature/mmmplotsuite-arviz

1dbe00b

williambdean reviewed Nov 26, 2025

View reviewed changes

isofer changed the title ~~Feature/mmmplotsuite arviz~~ MMM Plot Suite Migration: arviz_plots Multi-Backend Support Nov 27, 2025

isofer added 5 commits November 27, 2025 17:33

addressed 2 comments

93eeb63

remove versionadded

cfea8e5

remove tests

3c82941

update docstrings

153583e

changes name of MMMConfig to MMMPlotConfig

00c269d

MMM Plot Suite Migration: arviz_plots Multi-Backend Support #2098

Are you sure you want to change the base?

MMM Plot Suite Migration: arviz_plots Multi-Backend Support #2098

Conversation

isofer commented Nov 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

MMM Plot Suite Migration: arviz_plots Multi-Backend Support

Summary

Key Changes

🎨 New Multi-Backend Plotting Suite (plot.py)

🔄 Backward Compatibility & Legacy Suite

⚙️ Configuration System (config.py - NEW)

📊 API Improvements

🧪 Comprehensive Testing (~3,100 lines of test code)

Breaking Changes

Return Types

Parameter Changes

Method Changes

Migration Guide

Quick Start

Opting Out (Temporary)

Files Changed

New Files

Modified Files

Testing

Documentation

Timeline

Checklist

Notes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

codecov bot commented Nov 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

cetagostini left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI commented Nov 26, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

isofer commented Nov 27, 2025

Uh oh!

Reviewers

isofer commented Nov 21, 2025 •

edited

Loading

🎨 New Multi-Backend Plotting Suite (`plot.py`)

⚙️ Configuration System (`config.py` - NEW)

codecov bot commented Nov 22, 2025 •

edited

Loading