Add user customization to plot_curve methods (#1018)

Author: williambdean

pymc_marketing/mmm/components/base.py

@@ -21,19 +21,23 @@
 """
 
 import warnings
+from collections.abc import Iterable
 from copy import deepcopy
 from inspect import signature
 from typing import Any
 
-import matplotlib.pyplot as plt
 import numpy as np
 import numpy.typing as npt
 import pymc as pm
 import xarray as xr
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
 from pymc.distributions.shape_utils import Dims
 from pytensor import tensor as pt
+from pytensor.tensor.variable import TensorVariable
 
 from pymc_marketing.mmm.plot import (
+    SelToString,
     plot_curve,
     plot_hdi,
     plot_samples,
@@ -299,12 +303,10 @@ def variable_mapping(self) -> dict[str, str]:
 
     def _create_distributions(
         self, dims: Dims | None = None
-    ) -> dict[str, pt.TensorVariable]:
+    ) -> dict[str, TensorVariable]:
         dim_handler: DimHandler = create_dim_handler(dims)
 
-        def create_variable(
-            parameter_name: str, variable_name: str
-        ) -> pt.TensorVariable:
+        def create_variable(parameter_name: str, variable_name: str) -> TensorVariable:
             dist = self.function_priors[parameter_name]
             var = dist.create_variable(variable_name)
             return dim_handler(var, dist.dims)
@@ -344,7 +346,12 @@ def plot_curve(
         subplot_kwargs: dict | None = None,
         sample_kwargs: dict | None = None,
         hdi_kwargs: dict | None = None,
-    ) -> tuple[plt.Figure, npt.NDArray[plt.Axes]]:
+        axes: npt.NDArray[Axes] | None = None,
+        same_axes: bool = False,
+        colors: Iterable[str] | None = None,
+        legend: bool | None = None,
+        sel_to_string: SelToString | None = None,
+    ) -> tuple[Figure, npt.NDArray[Axes]]:
         """Plot curve HDI and samples.
 
         Parameters
@@ -357,6 +364,16 @@ def plot_curve(
             Keyword arguments for the plot_curve_sample function. Defaults to None.
         hdi_kwargs : dict, optional
             Keyword arguments for the plot_curve_hdi function. Defaults to None.
+        axes : npt.NDArray[plt.Axes], optional
+            The exact axes to plot on. Overrides any subplot_kwargs
+        same_axes : bool, optional
+            If the axes should be the same for all plots. Defaults to False.
+        colors : Iterable[str], optional
+            The colors to use for the plot. Defaults to None.
+        legend : bool, optional
+            If the legend should be shown. Defaults to None.
+        sel_to_string : SelToString, optional
+            The function to convert the selection to a string. Defaults to None.
 
         Returns
         -------
@@ -369,6 +386,11 @@ def plot_curve(
             subplot_kwargs=subplot_kwargs,
             sample_kwargs=sample_kwargs,
             hdi_kwargs=hdi_kwargs,
+            axes=axes,
+            same_axes=same_axes,
+            colors=colors,
+            legend=legend,
+            sel_to_string=sel_to_string,
         )
 
     def _sample_curve(
@@ -424,8 +446,8 @@ def plot_curve_samples(
         rng: np.random.Generator | None = None,
         plot_kwargs: dict | None = None,
         subplot_kwargs: dict | None = None,
-        axes: npt.NDArray[plt.Axes] | None = None,
-    ) -> tuple[plt.Figure, npt.NDArray[plt.Axes]]:
+        axes: npt.NDArray[Axes] | None = None,
+    ) -> tuple[Figure, npt.NDArray[Axes]]:
         """Plot samples from the curve.
 
         Parameters
@@ -466,8 +488,8 @@ def plot_curve_hdi(
         hdi_kwargs: dict | None = None,
         plot_kwargs: dict | None = None,
         subplot_kwargs: dict | None = None,
-        axes: npt.NDArray[plt.Axes] | None = None,
-    ) -> tuple[plt.Figure, npt.NDArray[plt.Axes]]:
+        axes: npt.NDArray[Axes] | None = None,
+    ) -> tuple[Figure, npt.NDArray[Axes]]:
         """Plot the HDI of the curve.
 
         Parameters
@@ -494,9 +516,10 @@ def plot_curve_hdi(
             axes=axes,
             subplot_kwargs=subplot_kwargs,
             plot_kwargs=plot_kwargs,
+            hdi_kwargs=hdi_kwargs,
         )
 
-    def apply(self, x: pt.TensorLike, dims: Dims | None = None) -> pt.TensorVariable:
+    def apply(self, x: pt.TensorLike, dims: Dims | None = None) -> TensorVariable:
         """Call within a model context.
 
         Used internally of the MMM to apply the transformation to the data.

pymc_marketing/mmm/fourier.py

@@ -205,7 +205,7 @@
 
 """
 
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 from typing import Any
 
 import arviz as az
@@ -219,7 +219,7 @@
 from typing_extensions import Self
 
 from pymc_marketing.constants import DAYS_IN_MONTH, DAYS_IN_YEAR
-from pymc_marketing.mmm.plot import plot_curve, plot_hdi, plot_samples
+from pymc_marketing.mmm.plot import SelToString, plot_curve, plot_hdi, plot_samples
 from pymc_marketing.prior import Prior, create_dim_handler
 
 X_NAME: str = "day"
@@ -465,6 +465,11 @@ def plot_curve(
         subplot_kwargs: dict | None = None,
         sample_kwargs: dict | None = None,
         hdi_kwargs: dict | None = None,
+        axes: npt.NDArray[plt.Axes] | None = None,
+        same_axes: bool = False,
+        colors: Iterable[str] | None = None,
+        legend: bool | None = None,
+        sel_to_string: SelToString | None = None,
     ) -> tuple[plt.Figure, npt.NDArray[plt.Axes]]:
         """Plot the seasonality for one full period.
 
@@ -478,6 +483,16 @@ def plot_curve(
             Keyword arguments for the plot_full_period_samples method, by default None
         hdi_kwargs : dict, optional
             Keyword arguments for the plot_full_period_hdi method, by default None
+        axes : npt.NDArray[plt.Axes], optional
+            Matplotlib axes, by default None
+        same_axes : bool, optional
+            Use the same axes for all plots, by default False
+        colors : Iterable[str], optional
+            Colors for the different plots, by default None
+        legend : bool, optional
+            Show the legend, by default None
+        sel_to_string : SelToString, optional
+            Function to convert the selection to a string, by default None
 
         Returns
         -------
@@ -491,6 +506,11 @@ def plot_curve(
             subplot_kwargs=subplot_kwargs,
             sample_kwargs=sample_kwargs,
             hdi_kwargs=hdi_kwargs,
+            axes=axes,
+            same_axes=same_axes,
+            colors=colors,
+            legend=legend,
+            sel_to_string=sel_to_string,
         )
 
     def plot_curve_hdi(
@@ -596,9 +616,9 @@ class YearlyFourier(FourierBase):
         dist = Prior("Laplace", mu=mu, b=b, dims="fourier")
         yearly = YearlyFourier(n_order=2, prior=dist)
         prior = yearly.sample_prior(random_seed=rng)
-        curve = yearly.sample_full_period(prior)
+        curve = yearly.sample_curve(prior)
 
-        _, axes = yearly.plot_full_period(curve)
+        _, axes = yearly.plot_curve(curve)
         axes[0].set(title="Yearly Fourier Seasonality")
         plt.show()
 
@@ -643,9 +663,9 @@ class MonthlyFourier(FourierBase):
         dist = Prior("Laplace", mu=mu, b=b, dims="fourier")
         yearly = MonthlyFourier(n_order=2, prior=dist)
         prior = yearly.sample_prior(samples=100)
-        curve = yearly.sample_full_period(prior)
+        curve = yearly.sample_curve(prior)
 
-        _, axes = yearly.plot_full_period(curve)
+        _, axes = yearly.plot_curve(curve)
         axes[0].set(title="Monthly Fourier Seasonality")
         plt.show()
 

pymc_marketing/mmm/linear_trend.py

@@ -52,19 +52,22 @@
 
 """
 
+from collections.abc import Iterable
 from typing import Any, cast
 
-import matplotlib.pyplot as plt
 import numpy as np
 import numpy.typing as npt
 import pymc as pm
 import pytensor.tensor as pt
 import xarray as xr
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
 from pydantic import BaseModel, Field, InstanceOf, model_validator
 from pymc.distributions.shape_utils import Dims
+from pytensor.tensor.variable import TensorVariable
 from typing_extensions import Self
 
-from pymc_marketing.mmm.plot import plot_curve
+from pymc_marketing.mmm.plot import SelToString, plot_curve
 from pymc_marketing.prior import Prior, create_dim_handler
 
 
@@ -278,7 +281,7 @@ def default_priors(self) -> dict[str, Prior]:
 
         return priors
 
-    def apply(self, t: pt.TensorLike) -> pt.TensorVariable:
+    def apply(self, t: pt.TensorLike) -> TensorVariable:
         """Create the linear trend for the given x values.
 
         Parameters
@@ -409,7 +412,12 @@ def plot_curve(
         sample_kwargs: dict | None = None,
         hdi_kwargs: dict | None = None,
         include_changepoints: bool = True,
-    ) -> tuple[plt.Figure, npt.NDArray[plt.Axes]]:
+        axes: npt.NDArray[Axes] | None = None,
+        same_axes: bool = False,
+        colors: Iterable[str] | None = None,
+        legend: bool | None = None,
+        sel_to_string: SelToString | None = None,
+    ) -> tuple[Figure, npt.NDArray[Axes]]:
         """Plot the curve samples from the trend.
 
         Parameters
@@ -424,6 +432,16 @@ def plot_curve(
             Keyword arguments for the HDI, by default None.
         include_changepoints : bool, optional
             Include the change points in the plot, by default True.
+        axes : npt.NDArray[plt.Axes], optional
+            Axes to plot the curve, by default None.
+        same_axes : bool, optional
+            Use the same axes for the samples, by default False.
+        colors : Iterable[str], optional
+            Colors for the samples, by default None.
+        legend : bool, optional
+            Include a legend in the plot, by default None.
+        sel_to_string : SelToString, optional
+            Function to convert the selection to a string, by default None.
 
         Returns
         -------
@@ -437,6 +455,11 @@ def plot_curve(
             subplot_kwargs=subplot_kwargs,
             sample_kwargs=sample_kwargs,
             hdi_kwargs=hdi_kwargs,
+            axes=axes,
+            same_axes=same_axes,
+            colors=colors,
+            legend=legend,
+            sel_to_string=sel_to_string,
         )
 
         if not include_changepoints: