Different media transformations across different channels

[wd60622]

Following #632, it is very easy to apply different media transformations to different subsets of channels.

Here is an example API:

Some functions in examples:

@dataclass
class MediaConfig:
    name: str
    columns: list[str]
    adstock: AdstockTransformation
    saturation: SaturationTransformation
    adstock_first: bool

def get_media_values(media_configs: list[MediaConfig]) -> list[str]: 
    """Get the different levels used in the configurations."""
    
def apply_media_transformations(
    media_data: pt.TensorLike, 
    media_configs: list[MediaConfig], 
    model: pm.Model,
) -> pt.TensorVariable:
    """Apply the media transformations to subsets according to media_configs."""

Then using it could look like this:

media_configs: list[MediaConfig]= [
    MediaConfig(
        name="offline",
        columns=["TV", "Radio"],
        adstock=GeometricAdstock(l_max=15),
        saturation=HillSaturation(),
        adstock_first=True,
    ),
    MediaConfig(
        name="online",
        columns=["Facebook", "Instagram", "YouTube", "TikTok"],
        adstock=GeometricAdstock(l_max=10),
        saturation=MichaelisMentenSaturation(),
        adstock_first=False,
    ),
]

media_values = get_media_values(media_configs)

coords = {
    "time": dates,
    "media": media_values,
}
model = pm.Model(coords=coords)

df: pd.DataFrame = ...  # Some data with all the columns

with model:
    media_data = pm.Data(
        "media_data",
        df.loc[:, media_values].to_numpy(),
        dims=("time", "media"),
    )
    media_transformation_data = pm.Deterministic(
        "channel_contributions",
        apply_media_transformation(media_data, media_configs, model),
        dims=("time", "media"),
    )

I would be interested in hearing if others are interested in this functionality.