🐼
Pandas DataFrame .assign の活用
pysparkやpolarsで.select内で演算をするような感覚でpandasを操作したい時があります。
# Iris データセット前処理 (pyspark)
df.select(
(F.col("petal length (cm)") * F.col("petal width (cm)")).alias("petal_area"),
(F.col("sepal length (cm)") * F.col("sepal width (cm)")).alias("sepal_area"),
...
)
# or
(
df
.withColumn("petal_area", F.col("petal length (cm)") * F.col("petal width (cm)"))
.withColumn("sepal_area", F.col("sepal length (cm)") * F.col("sepal width (cm)"))
)
Pandasでカラム作成するとき、以下のよう書けます。
df["petal_area"] = df["petal length (cm)"] * df["petal width (cm)"]
df["sepal_area"] = df["sepal length (cm)"] * df["sepal width (cm)"]
ただ、処理途中でエラーになる等すると、データフレームが中途半端に処理されてしまうため、操作を間違えるとデータの加工の再現性がなくなったり面倒です。
pandasは.assignを使うと便利です。
# irisデータセット
import pandas as pd
import sklearn
data = sklearn.datasets.load_iris()
data_df = pd.DataFrame(data.data, columns=data.feature_names)
y = pd.Series(data.target, name='target')
data_df.head(5)
╭────┬─────────────────────┬────────────────────┬─────────────────────┬────────────────────╮
│ │ sepal length (cm) │ sepal width (cm) │ petal length (cm) │ petal width (cm) │
├────┼─────────────────────┼────────────────────┼─────────────────────┼────────────────────┤
│ 0 │ 5.1 │ 3.5 │ 1.4 │ 0.2 │
├────┼─────────────────────┼────────────────────┼─────────────────────┼────────────────────┤
│ 1 │ 4.9 │ 3 │ 1.4 │ 0.2 │
├────┼─────────────────────┼────────────────────┼─────────────────────┼────────────────────┤
│ 2 │ 4.7 │ 3.2 │ 1.3 │ 0.2 │
├────┼─────────────────────┼────────────────────┼─────────────────────┼────────────────────┤
│ 3 │ 4.6 │ 3.1 │ 1.5 │ 0.2 │
├────┼─────────────────────┼────────────────────┼─────────────────────┼────────────────────┤
│ 4 │ 5 │ 3.6 │ 1.4 │ 0.2 │
╰────┴─────────────────────┴────────────────────┴─────────────────────┴────────────────────╯
# 特徴量エンジニアリング（例）
X = (
data_df
.assign(
# 先ほどの例
petal_area = lambda x: x["petal length (cm)"] * x["petal width (cm)"],
sepal_area = lambda x: x["sepal length (cm)"] * x["sepal width (cm)"],
# さらに追加
petal_length_to_width_ratio = lambda x: x["petal length (cm)"] / x["petal width (cm)"],
sepal_length_to_width_ratio = lambda x: x["sepal length (cm)"] / x["sepal width (cm)"],
# 同一assign内で作成した列を使える
is_petal_larger_than_p95 = lambda x: x["petal_area"] > x["petal_area"].quantile(0.95),
is_sepal_larger_than_p95 = lambda x: x["sepal_area"] > x["sepal_area"].quantile(0.95),
petal_to_sepal_area_ratio = lambda x: x["petal_area"] / x["sepal_area"],
# グループ分け・その内における平均値との差の計算の例
size_category = lambda x: pd.cut(
x["petal_area"],
bins=[-float("inf"), x["petal_area"].quantile(0.33), x["petal_area"].quantile(0.66), float("inf")],
labels=[0, 1, 2], # 小→中→大
),
mean_petal_area_in_size_group = lambda x: x.groupby("size_category", observed=True)["petal_area"].transform("mean"),
diff_from_group_mean = lambda x: x["petal_area"] - x["mean_petal_area_in_size_group"],
)
)
X.head(5)
╭────┬─────────────────────┬────────────────────┬─────────────────────┬────────────────────┬──────────────┬──────────────┬───────────────────────────────┬───────────────────────────────┬────────────────────────────┬────────────────────────────┬─────────────────────────────┬─────────────────┬─────────────────────────────────┬────────────────────────╮
│ │ sepal length (cm) │ sepal width (cm) │ petal length (cm) │ petal width (cm) │ petal_area │ sepal_area │ petal_length_to_width_ratio │ sepal_length_to_width_ratio │ is_petal_larger_than_p95 │ is_sepal_larger_than_p95 │ petal_to_sepal_area_ratio │ size_category │ mean_petal_area_in_size_group │ diff_from_group_mean │
├────┼─────────────────────┼────────────────────┼─────────────────────┼────────────────────┼──────────────┼──────────────┼───────────────────────────────┼───────────────────────────────┼────────────────────────────┼────────────────────────────┼─────────────────────────────┼─────────────────┼─────────────────────────────────┼────────────────────────┤
│ 0 │ 5.1 │ 3.5 │ 1.4 │ 0.2 │ 0.28 │ 17.85 │ 7 │ 1.45714 │ False │ False │ 0.0156863 │ 0 │ 0.3656 │ -0.0856 │
├────┼─────────────────────┼────────────────────┼─────────────────────┼────────────────────┼──────────────┼──────────────┼───────────────────────────────┼───────────────────────────────┼────────────────────────────┼────────────────────────────┼─────────────────────────────┼─────────────────┼─────────────────────────────────┼────────────────────────┤
│ 1 │ 4.9 │ 3 │ 1.4 │ 0.2 │ 0.28 │ 14.7 │ 7 │ 1.63333 │ False │ False │ 0.0190476 │ 0 │ 0.3656 │ -0.0856 │
├────┼─────────────────────┼────────────────────┼─────────────────────┼────────────────────┼──────────────┼──────────────┼───────────────────────────────┼───────────────────────────────┼────────────────────────────┼────────────────────────────┼─────────────────────────────┼─────────────────┼─────────────────────────────────┼────────────────────────┤
│ 2 │ 4.7 │ 3.2 │ 1.3 │ 0.2 │ 0.26 │ 15.04 │ 6.5 │ 1.46875 │ False │ False │ 0.0172872 │ 0 │ 0.3656 │ -0.1056 │
├────┼─────────────────────┼────────────────────┼─────────────────────┼────────────────────┼──────────────┼──────────────┼───────────────────────────────┼───────────────────────────────┼────────────────────────────┼────────────────────────────┼─────────────────────────────┼─────────────────┼─────────────────────────────────┼────────────────────────┤
│ 3 │ 4.6 │ 3.1 │ 1.5 │ 0.2 │ 0.3 │ 14.26 │ 7.5 │ 1.48387 │ False │ False │ 0.0210379 │ 0 │ 0.3656 │ -0.0656 │
├────┼─────────────────────┼────────────────────┼─────────────────────┼────────────────────┼──────────────┼──────────────┼───────────────────────────────┼───────────────────────────────┼────────────────────────────┼────────────────────────────┼─────────────────────────────┼─────────────────┼─────────────────────────────────┼────────────────────────┤
│ 4 │ 5 │ 3.6 │ 1.4 │ 0.2 │ 0.28 │ 18 │ 7 │ 1.38889 │ False │ False │ 0.0155556 │ 0 │ 0.3656 │ -0.0856 │
╰────┴─────────────────────┴────────────────────┴─────────────────────┴────────────────────┴──────────────┴──────────────┴───────────────────────────────┴───────────────────────────────┴────────────────────────────┴────────────────────────────┴─────────────────────────────┴─────────────────┴─────────────────────────────────┴────────────────────────╯
# 分類モデルで学習
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)
accuracy = model.score(X_test, y_test)
print(f"Model accuracy: {accuracy:.2f}")
Model accuracy: 1.00
Pro:
- .assign内部でエラーがあると.assign全体が処理されないので、中途半端な処理がデータフレームに適用されません。
- 好みにもよりますが、各行に処理をまとめられるので、可読性が上がります
Con:
- pandasの設定次第では.assign内ではデータフレームがコピーされ、メモリ使用量が増えるリスクがありますが、経験上そこまで困ったことはありません。相当の大容量サイズのデータフレームを扱うときだけ気をつければ良い（ただしその場合はsparkとか他に適当なツールがある）と思います。
Discussion