docs/custom_priors.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8hofhH9sdpFO"
      },
      "source": [
        "# Custom priors in LightweigthMMM"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xtBHFhN-sIo5"
      },
      "source": [
        "LightweightMMM allows you to pass your own prior to any of the parameters of the\n",
        "model. This notebook explains how this can be done."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8bSzzQqMpW2W"
      },
      "source": [
        "The `media_prior` has a dedicated parameter and is not part of the custom priors\n",
        "as it is a required argument and not an optional one. In this notebook we focus\n",
        "on the optional custom priors for the rest of the parameters in the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p0CRDUJUsINo"
      },
      "outputs": [],
      "source": [
        "# Please note that the values given here are just to demonstrate the usage of\n",
        "# the API provided and they are by no means intended to ressemble good prior`\n",
        "# values."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Udoqg9Iop0ji"
      },
      "source": [
        "### Initial setup. \n",
        "###### Refer to end to end examples for general information about workflow the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JQTOdErIeWlJ"
      },
      "outputs": [],
      "source": [
        "# Import jax.numpy and any other library we might need.\n",
        "import jax.numpy as jnp\n",
        "import numpyro"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mo29Y56CgHOl"
      },
      "outputs": [],
      "source": [
        "from lightweight_mmm import lightweight_mmm\n",
        "from lightweight_mmm import preprocessing\n",
        "from lightweight_mmm import utils"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SGX5QNHDecC_"
      },
      "outputs": [],
      "source": [
        "data_size = 104\n",
        "media_data, extra_features, target, costs = utils.simulate_dummy_data(\n",
        "    data_size=data_size + 13,\n",
        "    n_media_channels=3,\n",
        "    n_extra_features=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MBdYz2cGedww"
      },
      "outputs": [],
      "source": [
        "# Split and scale data.\n",
        "split_point = data_size - 13\n",
        "# Media data\n",
        "media_data_train = media_data[:split_point, ...]\n",
        "media_data_test = media_data[split_point:, ...]\n",
        "# Extra features\n",
        "extra_features_train = extra_features[:split_point, ...]\n",
        "extra_features_test = extra_features[split_point:, ...]\n",
        "# Target\n",
        "target_train = target[:split_point]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eikqG7w2e6kh"
      },
      "outputs": [],
      "source": [
        "media_scaler = preprocessing.CustomScaler(divide_operation=jnp.mean)\n",
        "extra_features_scaler = preprocessing.CustomScaler(divide_operation=jnp.mean)\n",
        "target_scaler = preprocessing.CustomScaler(divide_operation=jnp.mean)\n",
        "cost_scaler = preprocessing.CustomScaler(divide_operation=jnp.mean)\n",
        "\n",
        "media_data_train = media_scaler.fit_transform(media_data_train)\n",
        "extra_features_train = extra_features_scaler.fit_transform(extra_features_train)\n",
        "target_train = target_scaler.fit_transform(target_train)\n",
        "costs = cost_scaler.fit_transform(costs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PvDMbsTze_C-"
      },
      "outputs": [],
      "source": [
        "mmm = lightweight_mmm.LightweightMMM(model_name=\"carryover\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xRngSTPzfAki"
      },
      "outputs": [],
      "source": [
        "# Just a small number as in the following examples we are just seeing\n",
        "# how different parameters work in the `fit` method.\n",
        "number_warmup = 10\n",
        "number_samples = 10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WeQZ-H7txiAx"
      },
      "source": [
        "## What parameters does the model have? And what are their names?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MY4bRoNo1MBA"
      },
      "source": [
        "Although we go over all parameters here, for full detail on the model please refer to the [model documentation](https://lightweight-mmm.readthedocs.io/en/latest/models.html). The simplified model formulation is the following:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e5edKsZ-xtYE"
      },
      "source": [
        "$$ kpi_{t} = \\alpha + trend_{t} + seasonality_{t} + media\\_channels_{t} + other\\_factors_{t} $$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F31p5oOCyaJO"
      },
      "source": [
        "### Intercept:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eDv4O6Rkykaf"
      },
      "source": [
        "- $\\alpha \\sim HalfNormal(2)$\n",
        "- **Prior name: \"intercept\"**\n",
        "- Default prior: `numpyro.distributions.HalfNormal(scale=2)`\n",
        "- Final shape:\n",
        "  - National: `()`\n",
        "  - Geo: `(g,)` where `g` is the number of geos\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cEaT6rAh0xFr"
      },
      "source": [
        "### Trend:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n1vqMrF61cHA"
      },
      "source": [
        "- $trend_{t} = \\mu t^{\\kappa}$ \n",
        "- Where $t$ is a linear trend input\n",
        "- $\\mu \\sim Normal(0,1)$\n",
        "  - **Prior name: \"coef_trend\"**\n",
        "  - Default prior: `numpyro.distributions.Normal(loc=0., scale=1.)`\n",
        "  - Final shape:\n",
        "    - National: `()`\n",
        "    - Geo: `(g,)` where `g` is the number of geos\n",
        "- $\\kappa \\sim Uniform(0.5,1.5)$\n",
        "  - **Prior name: \"expo_trend\"**\n",
        "  - Default prior: `numpyro.distributions.Uniform(low=0.5, high=1.5)`\n",
        "  - Final shape:\n",
        "    - National: `()`\n",
        "    - Geo: `()`\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BqmKcyes1bar"
      },
      "source": [
        "### Seasonality:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lrOlay6pVZgq"
      },
      "source": [
        "#### Seasonality (for models using* **weekly observations**)\n",
        "- $seasonality_{t} = \\displaystyle\\sum_{d=1}^{2} (\\gamma_{1,d} cos(\\frac{2 \\pi d}{52}) + \\gamma_{2,d} sin(\\frac{2 \\pi d}{52}))$\n",
        "- $\\gamma_{1,d}, \\gamma_{2,d} \\sim Normal(0,1)$\n",
        "- **Prior name: \"gamma_seasonality\"**\n",
        "- Default prior: `numpyro.distributions.Normal(loc=0., scale=1.)`\n",
        "- Final shape:\n",
        "  - National: `(2, d)` where `d` is the number of degrees of seasonality\n",
        "  - Geo: `(2, d)` where `d` is the number of degrees of seasonality"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JHUjdlrYW7MP"
      },
      "source": [
        "#### Seasonality (for models using* **daily observations**)\n",
        "- $seasonality_{t} = \\displaystyle\\sum_{d=1}^{2} (\\gamma_{1,d} cos(\\frac{2 \\pi d}{365}) + \\gamma_{2,d} sin(\\frac{2 \\pi d}{365})) + \\displaystyle\\sum_{i=1}^{7} \\delta_{i}$\n",
        "- $\\gamma_{1,d}, \\gamma_{2,d} \\sim Normal(0,1)$\n",
        "  - **Prior name: \"gamma_seasonality\"**\n",
        "  - Default prior: `numpyro.distributions.Normal(loc=0., scale=1.)`\n",
        "  - Final shape:\n",
        "    - National: `(2, d)` where `d` is the number of degrees of seasonality\n",
        "    - Geo: `(2, d)` where `d` is the number of degrees of seasonality\n",
        "- $\\delta_{i} \\sim Normal(0,0.5)$\n",
        "  - **Prior name: \"weekday\"**\n",
        "  - Default prior: `numpyro.distributions.Normal(loc=0., scale=0.5)`\n",
        "  - Final shape:\n",
        "    - National: `(7,)`\n",
        "    - Geo: `(7,)`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ux0RN5nnItZ5"
      },
      "source": [
        "### Other Factors (extra features)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TAF5crP0LBh6"
      },
      "source": [
        "- $other\\_factors_{t} = \\displaystyle\\sum_{i=1}^{N} \\lambda_{i}Z_{i}$\n",
        "- $\\lambda_{i}  \\sim  Normal(0,1)$\n",
        "- Where $Z_{i}$ are other factors and $N$ is the number of other factors.\n",
        "  - **Prior name: \"coef_extra_features\"**\n",
        "  - Default prior: `numpyro.distributions.Normal(loc=0., scale=1.)`\n",
        "  - Final shape:\n",
        "    - National: `(f,)` where `f` is the number of extra features\n",
        "    - Geo: `(f,)` where `f` is the number of extra features"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f63u528ZMZSs"
      },
      "source": [
        "### Geo model only priors:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oyXHRtwvMfga"
      },
      "source": [
        "The geo model (hierarchical model) has the following other parameters that the \n",
        "national one does not have:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iWlh3rRnMoq0"
      },
      "source": [
        "$ kpi = ... + \\ \\tau \\ seasonality_{t} \\ + ... $"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mvGlPXoMNJsq"
      },
      "source": [
        "- $\\tau  \\sim  HalfNormal(0.5)$\n",
        "- **Prior name: \"coef_seasonality\"**\n",
        "- Default prior: `numpyro.distributions.HalfNormal(scale=.5)`\n",
        "- Final shape:\n",
        "    - Geo: `(g,)` where `g` is the number of geos"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "si-tCPUeV9DI"
      },
      "source": [
        "### Other priors"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W7K2da2rV_XG"
      },
      "source": [
        "The target is \n",
        "\n",
        "- $target \\sim N(\\mu, \\sigma) $\n",
        "- $\\sigma \\sim Gamma(1, 1)$\n",
        "- **Prior name: \"sigma\"**\n",
        "- Default prior: `numpyro.distributions.Gamma(concentration=1., rate=1.)`\n",
        "- Final shape:\n",
        "    - National: `()`\n",
        "    - Geo: `(g,)` where `g` is the number of geos"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "txe1LVFsg-Yt"
      },
      "source": [
        "### Media transformation priors:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9rdfBqvwhscZ"
      },
      "source": [
        "### Saturation\n",
        "\n",
        "\n",
        "#### Hill:\n",
        "\n",
        "- $media\\ channels_{t} = \\frac{1}{1+(x_{t,m}^{*} / K_{m})^{-S_{m}}}$\n",
        "- $K_{m} \\sim Gamma(1,1)$ \n",
        "  - It should be stricly positive\n",
        "  - **Prior name: \"half_max_effective_concentration\"**\n",
        "  - Default prior: `numpyro.distributions.Gamma(concentration=1., rate=1.)`\n",
        "  - Final shape:\n",
        "    - National: `(c,)` where `c` is the number of media channels\n",
        "    - Geo: `(c,)` where `c` is the number of media channels\n",
        "- $S_{m} \\sim Gamma(1,1)$\n",
        "  - **Prior name: \"slope\"**\n",
        "  - Default prior: `numpyro.distributions.Gamma(concentration=1., rate=1.)`\n",
        "  - Final shape:\n",
        "    - National: `(c,)` where `c` is the number of media channels\n",
        "    - Geo: `(c,)` where `c` is the number of media channels\n",
        "\n",
        "#### Exponent:\n",
        "\n",
        "\n",
        "- $media\\_channels_{t} = x_{t,m,s}^{*\\rho_{m}}$\n",
        "- $\\rho_{m} \\sim Beta(9,1)$\n",
        "- **Prior name: \"exponent\"**\n",
        "- Default prior: `numpyro.distributions.Beta(concentration1=9., concentration0=1.)`\n",
        "- Final shape:\n",
        "    - National: `(c,)` where `c` is the number of media channels\n",
        "    - Geo: `(c,)` where `c` is the number of media channels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SdTWBXOEhq7L"
      },
      "source": [
        "### Lagging\n",
        "\n",
        "#### Adstock\n",
        "\n",
        "- $media\\ channels_{t} =  x_{t,m} + \\lambda_{m} x_{t-1,m}^{*}$ where $t=2,..,N$\n",
        "- $\\lambda_{m} \\sim Beta(2,1)$\n",
        "- **Prior name: \"lag_weight\"**\n",
        "- Default prior: `numpyro.distributions.Beta(concentration1=2., concentration0=1.)`\n",
        "- Final shape:\n",
        "  - National: `(c,)` where `c` is the number of media channels\n",
        "  - Geo: `(c,)` where `c` is the number of media channels\n",
        "\n",
        "\n",
        "#### Carryover\n",
        "\n",
        "- $media\\ channels_{t} = \\frac{\\displaystyle\\sum_{l=0}^{L} \\tau_{m}^{(l-\\theta_{m})^2}x_{t-l,m}}{\\displaystyle\\sum_{l=0}^{L}\\tau_{m}^{(l-\\theta_{m})^2}}$\n",
        "- where $L=13$ for weekly data and $L=13*7$ for daily data\n",
        "- $\\tau_{m} \\sim Beta(1,1)$\n",
        "  - **Prior name: \"ad_effect_retention_rate\"**\n",
        "  - Default prior: `numpyro.distributions.Beta(concentration1=1., concentration0=1.)`\n",
        "  - Final shape:\n",
        "    - National: `(c,)` where `c` is the number of media channels\n",
        "    - Geo: `(c,)` where `c` is the number of media channels\n",
        "- $\\theta_{m} \\sim HalfNormal(2)$\n",
        "  - **Prior name: \"peak_effect_delay\"**\n",
        "  - Default prior: `numpyro.distributions.HalfNormal(scale=2.)`\n",
        "  - Final shape:\n",
        "    - National: `(c,)` where `c` is the number of media channels\n",
        "    - Geo: `(c,)` where `c` is the number of media channels\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8ulVJoLwxhUT"
      },
      "source": [
        "## What does the API to change them look like?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "haqwpXMwh23E"
      },
      "source": [
        "Media priors are always required to be given and they have a dedicated \n",
        "parameter for it in the `fit` method. Below we focus on the rest of the parameters of the model which are optional and are given through a dictionary\n",
        "with the prior name and its respective value/distribution."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VQSeTUIPfcDO"
      },
      "source": [
        "There are two main ways of passing your own priors to the parameter's models.\n",
        "- Passing a new distribution object.\n",
        "- Passing the values of the constructor for the default prior."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s50GUJLnfCzn"
      },
      "outputs": [],
      "source": [
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    custom_priors={\"intercept\": numpyro.distributions.HalfNormal(5)})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZP7TJ8XhQxhe"
      },
      "outputs": [],
      "source": [
        "# Also know that even if you pass the same prior as our default one you can \n",
        "# always do it from scratch and provide the whole object.\n",
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    custom_priors={\"intercept\": numpyro.distributions.HalfNormal(scale=2.)})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U_joNZjRpNqs"
      },
      "source": [
        "You can also find the full list of available distributions in the [Numpyro \n",
        "documentation](https://num.pyro.ai/en/stable/distributions.html)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7vktDhWLoGdh"
      },
      "source": [
        "Our default \"intercept\" prior is a `numpyro.distribution.HalfNormal(scale=2.)`. See [numpyro distribution here](https://num.pyro.ai/en/stable/distributions.html#halfnormal).\n",
        "\n",
        "Since the HalfNormal distribution has one positional argument (not positional only) we have the following ways you can give your desired values:\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2xMCse5Qp_SK"
      },
      "source": [
        "### As keyword arguments:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uTzWVQ0UhlE1"
      },
      "outputs": [],
      "source": [
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    custom_priors={\"intercept\": {\"scale\": 4.}})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FxOTQ8Osppnh"
      },
      "source": [
        "### As positional arguments:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9CDVhYdppdzn"
      },
      "outputs": [],
      "source": [
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    custom_priors={\"intercept\": (4.,)})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2E4BEqQ3ps5e"
      },
      "outputs": [],
      "source": [
        "# Another way of giving the first parameter (in case it has multiple  ones).\n",
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    custom_priors={\"intercept\": 2.})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HaCCmBJfqfk-"
      },
      "source": [
        "### When Numpyro distribution does not have default values"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q-2-1-bLqrpm"
      },
      "source": [
        "For example the [Beta](https://num.pyro.ai/en/stable/distributions.html#beta)\n",
        "distribution does not have default parameters so we will have to pass both \n",
        "parameters for altering our default priors.\n",
        "\n",
        "The default prior distribution of the exponent is a Beta \n",
        "distribution, let's see how it would work:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6z-YQBKCqk0x"
      },
      "outputs": [],
      "source": [
        "# The following will fail:\n",
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    custom_priors={\"exponent\": 0.65})\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eJf-YFc0uQyl"
      },
      "source": [
        "We need to specify all arguments either positionaly or by keyword:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oECeXK-jsD_h"
      },
      "outputs": [],
      "source": [
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    custom_priors={\"exponent\": (0.5, 1.5)})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5wg1cXZ3sRpY"
      },
      "outputs": [],
      "source": [
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    custom_priors={\"exponent\": {\"concentration1\": 0.7, \"concentration0\": 1.7}})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_yzc1bIesZRm"
      },
      "source": [
        "## Final shape of priors"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bCHEtV74MMD6"
      },
      "source": [
        "Prior values are in many cases arrays and not single values. We provide \n",
        "information about the target shape of each parameter so you can tailor your\n",
        "prior to that shape if you desired.\n",
        "\n",
        "Providing a single value will just broadcast the given value to the target \n",
        "shape."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "peU4j-THQYUF"
      },
      "source": [
        "For example our \"weekday\" prior has a target shape of (7,) let's see what the \n",
        "options are for passing a custom prior:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gcBOQ4TqQnBx"
      },
      "outputs": [],
      "source": [
        "# The simple case is where we provide single values as we have seen before.\n",
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    custom_priors={\"weekday\": numpyro.distributions.Normal(loc=0., scale=.5)})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wyc3dGosRhk0"
      },
      "outputs": [],
      "source": [
        "# But we can provide a shape that is broadcastable to (7,)\n",
        "weekday_prior = numpyro.distributions.Normal(loc=jnp.arange(7), scale=.5)\n",
        "\n",
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    custom_priors={\"weekday\": weekday_prior})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "newe33DfRycI"
      },
      "outputs": [],
      "source": [
        "# For one or all the parameters:\n",
        "weekday_prior = numpyro.distributions.Normal(\n",
        "    loc=jnp.arange(7), scale=jnp.array([2, 2, 2, 2, 2, 4, 4]))\n",
        "\n",
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    custom_priors={\"weekday\": weekday_prior})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UuinLG2FSDha"
      },
      "outputs": [],
      "source": [
        "# And also for just passing the values without the distribution object.\n",
        "\n",
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    custom_priors={\"weekday\": {\"loc\": jnp.arange(7), \"scale\": jnp.arange(7)}})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gf3wDlitQ2-i"
      },
      "outputs": [],
      "source": [
        "# And also for just passing the values as positional args.\n",
        "\n",
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    custom_priors={\"weekday\": (jnp.arange(7), jnp.arange(7))})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xDdK32ZYSfMH"
      },
      "source": [
        "The only exception will be when you only pass one value which is a Sequence\n",
        "as a positional argument and it is not encapsulated within another Sequence. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WjoxZl5aTBMf"
      },
      "outputs": [],
      "source": [
        "# This does not work\n",
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    # This does not work since we tried to pass the different values of the \n",
        "    # sequence to the prior distribution constructor which in this case only \n",
        "    # has 2 positional arguments.\n",
        "    custom_priors={\"weekday\": jnp.arange(7)})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tGsx28mHXCEl"
      },
      "outputs": [],
      "source": [
        "mmm.fit(\n",
        "    media=media_data_train,\n",
        "    media_prior=costs,\n",
        "    target=target_train,\n",
        "    extra_features=extra_features_train,\n",
        "    number_warmup=number_warmup,\n",
        "    number_samples=number_samples,\n",
        "    # This works since the prior is within the tuple\n",
        "    custom_priors={\"weekday\": (jnp.arange(7),)})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6bscWUOPYcH-"
      },
      "source": [
        "### Final notes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ELmnk_tRYdn8"
      },
      "source": [
        "To avoid silent errors or unexpected behaviour we also raise exceptions\n",
        "(in most casdes) when user passes a prior for a key that does not exist in the model."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "name": "custom_priors.ipynb",
      "private_outputs": true,
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}