From e6f844f278452ec49aa76d3e4061b4e912f38e84 Mon Sep 17 00:00:00 2001
From: Juan Orduz <juanitorduz@gmail.com>
Date: Mon, 22 Jul 2024 14:18:06 +0200
Subject: [PATCH] Various MMM small documentation fixes (#854)

* make nb ruff https://github.com/astral-sh/ruff-vscode/issues/546

* fix docs

* fix quickstart

* add types
---
 .pre-commit-config.yaml                     |  6 +++---
 docs/source/index.md                        | 16 +++++++++-----
 pymc_marketing/mmm/delayed_saturated_mmm.py | 24 +++++++++++----------
 pyproject.toml                              |  3 +++
 4 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 05af4713c..e1d9ee8f1 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -11,14 +11,14 @@ repos:
           - --exclude=docs/
           - --exclude=scripts/
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.5.3
+    rev: v0.5.4
     hooks:
       - id: ruff
-        types_or: [ python, pyi, jupyter ]
+        types_or: [python, pyi, jupyter]
         args: ["--fix", "--output-format=full"]
         exclude: ^docs/source/notebooks/clv/dev/
       - id: ruff-format
-        types_or: [ python, pyi, jupyter ]
+        types_or: [python, pyi, jupyter]
         exclude: ^docs/source/notebooks/clv/dev/
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.10.1
diff --git a/docs/source/index.md b/docs/source/index.md
index 538991a9e..58581d026 100644
--- a/docs/source/index.md
+++ b/docs/source/index.md
@@ -111,13 +111,20 @@ Start VS Code and ensure that the "Jupyter" extension is installed. Press Ctrl +
 
 ```python
 import pandas as pd
-from pymc_marketing.mmm import DelayedSaturatedMMM
+
+from pymc_marketing.mmm import (
+    GeometricAdstock,
+    LogisticSaturation,
+    MMM,
+)
 
 
 data_url = "https://raw.githubusercontent.com/pymc-labs/pymc-marketing/main/data/mmm_example.csv"
 data = pd.read_csv(data_url, parse_dates=['date_week'])
 
-mmm = DelayedSaturatedMMM(
+mmm = MMM(
+    adstock=GeometricAdstock(l_max=8),
+    saturation=LogisticSaturation(),
     date_column="date_week",
     channel_columns=["x1", "x2"],
     control_columns=[
@@ -125,7 +132,6 @@ mmm = DelayedSaturatedMMM(
         "event_2",
         "t",
     ],
-    adstock_max_lag=8,
     yearly_seasonality=2,
 )
 ```
@@ -133,8 +139,8 @@ mmm = DelayedSaturatedMMM(
 Initiate fitting and get a visualization of some of the outputs with:
 
 ```python
-X = data.drop('y',axis=1)
-y = data['y']
+X = data.drop("y",axis=1)
+y = data["y"]
 mmm.fit(X,y)
 mmm.plot_components_contributions();
 ```
diff --git a/pymc_marketing/mmm/delayed_saturated_mmm.py b/pymc_marketing/mmm/delayed_saturated_mmm.py
index 5db0b6667..eb39cf38c 100644
--- a/pymc_marketing/mmm/delayed_saturated_mmm.py
+++ b/pymc_marketing/mmm/delayed_saturated_mmm.py
@@ -1683,20 +1683,22 @@ def sample_posterior_predictive(
         Sample from the model's posterior predictive distribution.
 
         Parameters
-        ---------
+        ----------
         X_pred : array, shape (n_pred, n_features)
             The input data used for prediction.
-        extend_idata : Boolean determining whether the predictions should be added to inference data object.
+        extend_idata : bool, optional
+            Boolean determining whether the predictions should be added to inference data object. Defaults to True.
+        combined: bool, optional
+            Combine chain and draw dims into sample. Won't work if a dim named sample already exists. Defaults to True.
+        include_last_observations: bool, optional
+            Boolean determining whether to include the last observations of the training data in order to carry over
+            costs with the adstock transformation. Assumes that X_pred are the next predictions following the
+            training data.Defaults to False.
+        original_scale: bool, optional
+            Boolean determining whether to return the predictions in the original scale of the target variable.
             Defaults to True.
-        combined: Combine chain and draw dims into sample. Won't work if a dim named sample already exists.
-            Defaults to True.
-        include_last_observations: Boolean determining whether to include the last observations of the training
-            data in order to carry over costs with the adstock transformation.
-            Assumes that X_pred are the next predictions following the training data.
-            Defaults to False.
-        original_scale: Boolean determining whether to return the predictions in the original scale
-            of the target variable. Defaults to True.
-        **sample_posterior_predictive_kwargs: Additional arguments to pass to pymc.sample_posterior_predictive
+        **sample_posterior_predictive_kwargs
+            Additional arguments to pass to pymc.sample_posterior_predictive
 
         Returns
         -------
diff --git a/pyproject.toml b/pyproject.toml
index a2f3da1c3..6dec3b142 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -112,6 +112,9 @@ ignore = [
 [tool.ruff.lint.pycodestyle]
 max-line-length = 120
 
+[tool.ruff]
+extend-include = ["*.ipynb"]
+
 [tool.pytest.ini_options]
 addopts = [
     "-v",