polars.Exprを使った特徴量管理

polarsではpl.Exprを用いて、「カラムの変換操作」自体をオブジェクトとして扱うことができる。この機能が機械学習の特徴量管理と相性がいいのでは？という話。

実装

簡易的なコードを示す。

from collections import UserDict
from typing import Self

import polars as pl

class FeatureExpressions(UserDict):
    """
    特徴量表現とそのラベルを紐づけて格納するカスタムコレクション
    """

    def __setitem__(self, key: str, value: list[pl.Expr]) -> None:
        if not isinstance(key, str):
            raise TypeError("Key must be a string")
        if not isinstance(value, list) or not all(isinstance(v, pl.Expr) for v in value):
            raise TypeError("Value must be a list of polars expressions")
        self.data[key] = value

    def __getitem__(self, key: str) -> list[pl.Expr]:
        return list(super().__getitem__(key))

    def filter(self, feature_names: list[str]) -> Self:
        """
        特徴量をフィルタリングする
        """
        return FeatureExpressions({k: self.data[k] for k in feature_names if k in self.data})

def create_features(df: pl.DataFrame, feature_expressions: FeatureExpressions) -> pl.DataFrame:
    """
    与えられた特徴量表現に従って特徴量を加工する
    """
    all_expressions = []

    for feature_name, expressions in feature_expressions.items():
        for expr in expressions:
            expr_name = f"{feature_name}_{expr.meta.output_name()}"
            all_expressions.append(expr.alias(expr_name))

    return df.with_columns(all_expressions)

FeatureExpressionはpl.Exprの形で記述された特徴量を記録するマスターのようなもので、これを利用して多様な特徴量を管理する。実際の加工処理では、filterで加工対象の特徴量に絞ったものをcreate_featureに渡して所望の特徴量を作成する。

使用例

Kaggleコンペでの使用例。feature_master.py上で様々な特徴量をpl.Exprの形で定義する。

# feature_master.py
feature_expressions_master = FeatureExpressions()

feature_expressions_master["agent_property"] = [
    pl.col("agent1").str.extract(AGENT_PATTERN, 1).alias("p1_selection"),
    pl.col("agent1").str.extract(AGENT_PATTERN, 2).alias("p1_exploration").cast(pl.Float32),
    pl.col("agent1").str.extract(AGENT_PATTERN, 3).alias("p1_playout"),
    pl.col("agent1").str.extract(AGENT_PATTERN, 4).alias("p1_bounds"),
    pl.col("agent2").str.extract(AGENT_PATTERN, 1).alias("p2_selection"),
    pl.col("agent2").str.extract(AGENT_PATTERN, 2).alias("p2_exploration").cast(pl.Float32),
    pl.col("agent2").str.extract(AGENT_PATTERN, 3).alias("p2_playout"),
    pl.col("agent2").str.extract(AGENT_PATTERN, 4).alias("p2_bounds"),
]

feature_expressions_master["lud_rules"] = [
    pl.col("LudRules").str.extract(LUD_RULES_PATTERN, 1).alias("LudRules_game"),
    pl.col("LudRules").str.extract(LUD_RULES_PATTERN, 2).alias("LudRules_players"),
    pl.col("LudRules").str.extract(LUD_RULES_PATTERN, 3).alias("LudRules_equipment"),
    pl.col("LudRules").str.extract(LUD_RULES_PATTERN, 4).alias("LudRules_rules"),
]

feature_expressions_master["baseline_features"] = [
    (pl.col("NumRows") * pl.col("NumColumns")).alias("area"),
    (pl.col("NumColumns").eq(pl.col("NumRows"))).cast(pl.Int8).alias("row_equal_col"),
    (pl.col("PlayoutsPerSecond") / (pl.col("MovesPerSecond") + 1e-15)).alias("Playouts/Moves"),
    (pl.col("MovesPerSecond") / (pl.col("PlayoutsPerSecond") + 1e-15)).alias("EfficiencyPerPlayout"),
    (pl.col("DurationActions") / (pl.col("DurationTurnsStdDev") + 1e-15)).alias("TurnsDurationEfficiency"),
    (pl.col("AdvantageP1") / (pl.col("Balance") + 1e-15)).alias("AdvantageBalanceRatio"),
    (pl.col("DurationActions") / (pl.col("MovesPerSecond") + 1e-15)).alias("ActionTimeEfficiency"),
    (pl.col("DurationTurnsStdDev") / (pl.col("DurationActions") + 1e-15)).alias("StandardizedTurnsEfficiency"),
    (pl.col("AdvantageP1") / (pl.col("DurationActions") + 1e-15)).alias("AdvantageTimeImpact"),
    (pl.col("DurationActions") / (pl.col("StateTreeComplexity") + 1e-15)).alias("DurationToComplexityRatio"),
    (pl.col("GameTreeComplexity") / (pl.col("StateTreeComplexity") + 1e-15)).alias("NormalizedGameTreeComplexity"),
    (pl.col("Balance") * pl.col("GameTreeComplexity")).alias("ComplexityBalanceInteraction"),
    (pl.col("StateTreeComplexity") + pl.col("GameTreeComplexity")).alias("OverallComplexity"),
    (pl.col("GameTreeComplexity") / (pl.col("PlayoutsPerSecond") + 1e-15)).alias("ComplexityPerPlayout"),
    (pl.col("DurationTurnsNotTimeouts") / (pl.col("MovesPerSecond") + 1e-15)).alias("TurnsNotTimeouts/Moves"),
    (pl.col("Timeouts") / (pl.col("DurationActions") + 1e-15)).alias("Timeouts/DurationActions"),
    (pl.col("OutcomeUniformity") / (pl.col("AdvantageP1") + 1e-15)).alias("OutcomeUniformity/AdvantageP1"),
    (pl.col("StepDecisionToEnemy") + pl.col("SlideDecisionToEnemy") + pl.col("HopDecisionMoreThanOne")).alias("ComplexDecisionRatio"),
    (pl.col("StepDecisionToEnemy") + pl.col("HopDecisionEnemyToEnemy") + pl.col("HopDecisionFriendToEnemy") + pl.col("SlideDecisionToEnemy")).alias(
        "AggressiveActionsRatio"
    ),
]

...

実行スクリプト上で利用する特徴量を指定して加工処理を行う。

# run.py
use_features = ["agent_property", "lud_rules_features"]
feature_expressions = feature_expressions_master.filter(use_features)
df_result = create_feature(df, feature_expressions)

所感

実際の計算と切り離された形で「論理的に」特徴量管理ができるところがポイント。特徴量の論理的な表現を一か所にまとめることで見通しよく管理でき、かつ必要に応じて使いたい特徴量をピックアップして利用できる。

効率的に特徴量エンジニアリングの試行錯誤を回したり、チームで特徴量の定義を共有・管理したいときに使えるかもしれない。