feat: Add Yeo-Johnson transformer

[metaboulie]

@transformer
def yeojohnson(brack: tuple = (-2, 2)):
    """Applies the Yeo-Johnson transformation to numeric columns in a panel DataFrame.

    Parameters
    ----------
    brack : 2-tuple, optional
        The starting interval for a downhill bracket search with optimize.brent. Note that this
        is in most cases not critical; the final result is allowed to be outside this bracket.
    """

    def transform(X: pl.LazyFrame) -> pl.LazyFrame:
        idx_cols = X.columns[:2]
        entity_col, time_col = idx_cols
        gb = X.groupby(X.columns[0])
        # Step 1. Compute optimal lambdas
        lmbds = gb.agg(
            PL_NUMERIC_COLS(entity_col, time_col)
            .apply(lambda x: yeojohnson_normmax(x, brack))
            .cast(pl.Float64())
            .suffix("__lmbd")
        )
        # Step 2. Transform
        cols = X.select(PL_NUMERIC_COLS(entity_col, time_col)).columns
        X_new = X.join(lmbds, on=entity_col, how="left").select(
            idx_cols
            + [
                pl.when((pl.col(col) >= 0) & (pl.col(f"{col}__lmbd") == 0))
                .then(pl.col(col).log1p())
                .when(pl.col(col) >= 0)
                .then(((pl.col(col) + 1) ** pl.col(f"{col}__lmbd") - 1)
                      / pl.col(f"{col}__lmbd"))
                .when((pl.col(col) < 0) & (pl.col(f"{col}__lmbd") == 2))
                .then(-pl.col(col).log1p())
                .otherwise(-((-pl.col(col) + 1) ** (2 - pl.col(f"{col}__lmbd")) - 1)
                           / (2 - pl.col(f"{col}__lmbd")))
                for col in cols
            ]
        )
        artifacts = {"X_new": X_new, "lmbds": lmbds}
        return artifacts

    def invert(state: ModelState, X: pl.LazyFrame) -> pl.LazyFrame:
        entity_col, time_col = X.columns[:2]
        lmbds = state.artifacts["lmbds"]
        cols = X.select(PL_NUMERIC_COLS(entity_col, time_col)).columns
        X_new = (
            X.join(lmbds, on=entity_col, how="left", suffix="__lmbd")
            .with_columns(
                [
                    pl.when((pl.col(col) >= 0) & (pl.col(f"{col}__lmbd") == 0))
                    .then((pl.col(col).exp()) - 1)
                    .when(pl.col(col) >= 0)
                    .then((pl.col(col) * pl.col(f"{col}__lmbd") + 1) ** (1 / pl.col(f"{col}__lmbd")) - 1)
                    .when((pl.col(col) < 0) & (pl.col(f"{col}__lmbd") == 2))
                    .then(1 - (-(pl.col(col)).exp()))
                    .otherwise(
                        (1 - (-(2 - pl.col(f"{col}__lmbd")) * pl.col(col) + 1) **
                         (1 / (2 - pl.col(f"{col}__lmbd")))
                         )
                    )
                    for col in cols
                ]
            )
            .select(X.columns)
        )
        return X_new

    return transform, invert

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
functime-docs	✅ Ready (Inspect)	Visit Preview	💬 Add feedback	Sep 6, 2023 9:17am

[topher-lo]

Hi @metaboulie, thanks for making a contribution! We recently created a discord channel for functime: https://discord.gg/jqDyuFa3. Will be great to meet you there 🤝

[metaboulie]

Hi @metaboulie, thanks for making a contribution! We recently created a discord channel for functime: https://discord.gg/jqDyuFa3. Will be great to meet you there 🤝

I'd like to join the channel, but it appears that this invitation is invalid, could you please offer me another link or something? Thanks a lot 😶‍🌫️

[topher-lo]

@metaboulie https://discord.gg/J486Jezh here's a new link!

[topher-lo]

If it's still not working, here's my email: chris@neocortexdb.com just send over your email or discord username and I can send an invite directly!

[topher-lo]

LGTM. @metaboulie congrats on your first PR here! 🎊