# Transformers (`transformers.py`)

This module provides various functions for transforming marketing data, primarily focusing on adstock (carryover effects) and saturation (diminishing returns) transformations commonly used in Marketing Mix Models (MMMs).

## Protocols

### `SaturationTransformation`

```python
class SaturationTransformation(Protocol):
    """Protocol defining the interface for saturation transformations."""
    @property
    def function(self) -> Callable[..., TensorVariable]:
        """The mathematical function implementing the saturation."""
        ...

    @property
    def variable_mapping(self) -> Dict[str, str]:
        """Mapping from saturation function parameter names to model variable names."""
        ...
```

A protocol defining the required structure for saturation transformation functions used within the model, particularly for lift testing integration. Classes implementing this protocol must provide:

-   `function`: A callable representing the saturation function itself (e.g., `logistic_saturation`).
-   `variable_mapping`: A dictionary mapping the parameter names expected by the `function` to the corresponding variable names used within the PyMC model.

## Enums

### `ConvMode`

```python
class ConvMode(str, Enum):
    After = "After"
    Before = "Before"
    Overlap = "Overlap"
```

Enum defining the different modes for applying 1D convolution in `batched_convolution`, affecting how boundaries are handled.

-   `After`: Standard adstock effect with trailing decay.
-   `Before`: Leading effect ("excitement" or "wow factor").
-   `Overlap`: Effect overlaps both preceding and succeeding elements.

### `WeibullType`

```python
class WeibullType(str, Enum):
    PDF = "PDF"
    CDF = "CDF"
```

Enum defining the type of Weibull distribution function to use for the `weibull_adstock` transformation.

-   `PDF`: Uses the Probability Density Function for weighting.
-   `CDF`: Uses the complementary Cumulative Distribution Function (1 - CDF) for weighting.

## Functions

### `batched_convolution`

```python
def batched_convolution(
    x: TensorLike,
    w: TensorLike,
    axis: int = 0,
    mode: Union[ConvMode, str] = ConvMode.After,
) -> TensorVariable:
```

Applies a 1D convolution in a vectorized way across multiple batch dimensions. This is the core function used by the adstock transformations.

**Parameters:**

-   `x` (`TensorLike`): The input array to convolve.
-   `w` (`TensorLike`): The convolution weights. The last axis determines the number of steps (lag).
-   `axis` (`int`, optional): The axis of `x` along which to apply the convolution. Defaults to `0`.
-   `mode` (`Union[ConvMode, str]`, optional): The convolution mode (`ConvMode.After`, `ConvMode.Before`, `ConvMode.Overlap`). Defaults to `ConvMode.After`.

**Returns:**

-   `TensorVariable`: The result of convolving `x` with `w`. Shape matches `x` (considering broadcasting with `w`).

*(See the docstring in the source code for a plot illustrating the different modes.)*

### `geometric_adstock`

```python
def geometric_adstock(
    x: TensorLike,
    alpha: float = 0.0,
    l_max: int = 12,
    normalize: bool = False,
    axis: int = 0
) -> TensorVariable:
```

Applies the geometric adstock transformation, assuming the peak advertising effect occurs at the time of exposure.

**Parameters:**

-   `x` (`TensorLike`): Input tensor (e.g., media spend/impressions).
-   `alpha` (`float`, optional): Retention rate (0 to 1). Higher values mean longer-lasting effects. Defaults to `0.0`.
-   `l_max` (`int`, optional): Maximum duration (lag) of the carryover effect. Defaults to `12`.
-   `normalize` (`bool`, optional): Whether to normalize the adstock weights so they sum to 1. Defaults to `False`.
-   `axis` (`int`, optional): Axis along which to apply the transformation. Defaults to `0`.

**Returns:**

-   `TensorVariable`: The transformed tensor with geometric adstock applied.

*(See the docstring in the source code for a plot illustrating the effect of `alpha` and `normalize`.)*

### `delayed_adstock`

```python
def delayed_adstock(
    x: TensorLike,
    alpha: float = 0.0,
    theta: int = 0,
    l_max: int = 12,
    normalize: bool = False,
    axis: int = 0,
) -> TensorLike:
```

Applies a delayed adstock transformation, allowing the peak effect to occur after the initial exposure time.

**Parameters:**

-   `x` (`TensorLike`): Input tensor.
-   `alpha` (`float`, optional): Retention rate (0 to 1). Defaults to `0.0`.
-   `theta` (`int`, optional): Delay of the peak effect (0 to `l_max` - 1). Defaults to `0`.
-   `l_max` (`int`, optional): Maximum duration of carryover effect. Defaults to `12`.
-   `normalize` (`bool`, optional): Whether to normalize the weights. Defaults to `False`.
-   `axis` (`int`, optional): Axis along which to apply the transformation. Defaults to `0`.

**Returns:**

-   `TensorLike`: The transformed tensor with delayed adstock applied.

*(See the docstring in the source code for a plot illustrating the effect of `alpha`, `theta`, and `normalize`.)*

### `weibull_adstock`

```python
def weibull_adstock(
    x: TensorLike,
    lam: float = 1.0,
    k: float = 1.0,
    l_max: int = 12,
    axis: int = 0,
    type: Union[WeibullType, str] = WeibullType.PDF,
) -> TensorLike:
```

Applies an adstock transformation using weights derived from the Weibull distribution (either PDF or CDF). Offers more flexibility than geometric or delayed adstock.

**Parameters:**

-   `x` (`TensorLike`): Input tensor.
-   `lam` (`float`, optional): Scale parameter (lambda > 0) of the Weibull distribution. Defaults to `1.0`.
-   `k` (`float`, optional): Shape parameter (k > 0) of the Weibull distribution. Defaults to `1.0`.
-   `l_max` (`int`, optional): Maximum duration of carryover effect. Defaults to `12`.
-   `axis` (`int`, optional): Axis along which to apply the transformation. Defaults to `0`.
-   `type` (`Union[WeibullType, str]`, optional): Type of Weibull function to use (`WeibullType.PDF` or `WeibullType.CDF`). Defaults to `WeibullType.PDF`.

**Returns:**

-   `TensorLike`: The transformed tensor with Weibull adstock applied.

*(See the docstring in the source code for plots illustrating the effect of `lam`, `k`, and `type`.)*

### `logistic_saturation`

```python
def logistic_saturation(
    x: TensorLike,
    lam: Union[npt.NDArray[np.float_], float] = 0.5
) -> TensorVariable:
```

Applies a logistic saturation function to model diminishing returns.
Formula: `(1 - exp(-lam * x)) / (1 + exp(-lam * x))`

**Parameters:**

-   `x` (`TensorLike`): Input tensor (e.g., adstocked media spend).
-   `lam` (`Union[npt.NDArray[np.float_], float]`, optional): Saturation parameter(s). Higher values cause saturation to occur more quickly. Defaults to `0.5`.

**Returns:**

-   `TensorVariable`: The transformed tensor with logistic saturation applied.

*(See the docstring in the source code for a plot illustrating the effect of `lam`.)*

### `tanh_saturation`

```python
def tanh_saturation(
    x: TensorLike,
    b: TensorLike = 0.5,
    c: TensorLike = 0.5,
) -> TensorVariable:
```

Applies a hyperbolic tangent (tanh) saturation function.
Formula: `b * tanh(x / (b * c))`

**Parameters:**

-   `x` (`TensorLike`): Input tensor.
-   `b` (`TensorLike`, optional): Saturation level (maximum effect). Defaults to `0.5`.
-   `c` (`TensorLike`, optional): Controls the steepness / initial cost per acquisition. Defaults to `0.5`.

**Returns:**

-   `TensorVariable`: The transformed tensor with tanh saturation applied.

*(See the docstring in the source code for a plot illustrating the effect of `b` and `c`.)*

### `tanh_saturation_baselined`

```python
def tanh_saturation_baselined(
    x: TensorLike,
    x0: TensorLike,
    gain: TensorLike = 0.5,
    r: TensorLike = 0.5,
) -> TensorVariable:
```

Applies a reparameterized tanh saturation function based on a reference point (`x0`), gain at that point, and an overspend fraction. This parameterization can be more intuitive for setting priors based on domain knowledge.

**Parameters:**

-   `x` (`TensorLike`): Input tensor.
-   `x0` (`TensorLike`): Reference point (e.g., median spend) where gain and overspend fraction are defined.
-   `gain` (`TensorLike`, optional): The Return on Ad Spend (ROAS) or effectiveness at the reference point `x0`. Defined as `f(x0) / x0`. Defaults to `0.5`.
-   `r` (`TensorLike`, optional): The overspend fraction, indicating how close `f(x0)` is to the maximum saturation level `b`. Defined as `f(x0) / b`. Defaults to `0.5`.

**Returns:**

-   `TensorVariable`: The transformed tensor with baselined tanh saturation applied.

*(See the docstring in the source code for a plot illustrating the parameterization.)*

## Classes

### `LogisticSaturation`

```python
class LogisticSaturation:
    """Wrapper class for logistic_saturation to conform to SaturationTransformation protocol."""
    @property
    def function(self) -> Callable[[Any, Any], TensorVariable]: ...
    @property
    def variable_mapping(self) -> Dict[str, str]: ...
```

A simple wrapper class that makes the `logistic_saturation` function conform to the `SaturationTransformation` protocol. Used for integrating logistic saturation with features like lift testing.

-   `function`: Returns the `logistic_saturation` function.
-   `variable_mapping`: Returns `{"lam": "lam"}`, assuming the model variable for the saturation parameter is named `lam`.

## Named Tuples

### `TanhSaturationParameters`

```python
class TanhSaturationParameters(NamedTuple):
    b: TensorLike  # Saturation
    c: TensorLike  # Customer Acquisition Cost at 0
```

Represents the standard parameters (`b`: saturation, `c`: initial CAC) for the `tanh_saturation` function.

**Methods:**

-   `baseline(self, x0: TensorLike) -> TanhSaturationBaselinedParameters`: Converts these standard parameters to the baselined parameterization (`TanhSaturationBaselinedParameters`) given a reference point `x0`.

### `TanhSaturationBaselinedParameters`

```python
class TanhSaturationBaselinedParameters(NamedTuple):
    x0: TensorLike    # Baseline spend
    gain: TensorLike  # ROAS at x0
    r: TensorLike     # Overspend Fraction (f(x0) / saturation)
```

Represents the baselined parameters (`x0`: reference spend, `gain`: ROAS at `x0`, `r`: overspend fraction) for the `tanh_saturation_baselined` function.

**Methods:**

-   `debaseline(self) -> TanhSaturationParameters`: Converts these baselined parameters back to the standard parameterization (`TanhSaturationParameters`).
-   `rebaseline(self, x1: TensorLike) -> TanhSaturationBaselinedParameters`: Converts the parameters to be baselined at a *new* reference point `x1`.