Different media transformations across different channels

[williambdean]

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.

[williambdean]

@carlosagostini might have some opinions

Could start as notebook via #711

[juanitorduz]

Yes! This would be great! I think Pydanty can be helpful here ;)

[williambdean]

Yes! This would be great! I think Pydanty can be helpful here ;)

Agreed. I often use the from pydantic.dataclasses import dataclass but BaseModel could also be used

[williambdean]

Might be good to group the adstock, saturation, and adstock_first together. Could be something like this

@dataclass
class MediaTransformation: 
    adstock: AdstockTransformation
    saturation: SaturationTransformation
    adstock_first: bool = True

    def __post_init_(self) -> None:
        self.first, self.second = (
            (self.adstock, self.saturation) 
            if self.adstock_first 
            else (self.saturation, self.adstock)
        )

    def __call__(self, media_data: pt.TensorLike, dims: Dims) -> pt.TensorVariable:
        return self.second.apply(self.first.apply(media_data, dims), dims)

@dataclass
class MediaConfig: 
    name: str
    columns: list[str]
    transformation: MediaTransformation

This way the MediaTransformation can be used on its own as well. i.e. what is currently being called in the MMM class.