競馬AI開発メモ：Self-Attentionでレース内の相手関係を特徴量化してみる

# 01_make_race_dataset.py（データセット構築の核心部分）

import sqlite3, numpy as np, pandas as pd, json
from pathlib import Path

MAX_HORSES = 16
SPECIAL_FLAGS = ['失', '取', '除', '中', '取消レース', '失格レース', '除外レース', '中止レース']

# DB から全期間を読み込み（2017-2018 のデータも shift 計算に使う）
con = sqlite3.connect('common.db')
df = pd.read_sql('SELECT ... FROM result_table', con); con.close()

# 特殊フラグ・帯広 除外
mask = (df[SPECIAL_FLAGS].fillna(0).astype(int).sum(axis=1) == 0
        & df['着順'].between(1, 20)
        & (df['場所'] != '帯広'))
df = df[mask].sort_values(['馬名', '日付', 'レースID'])

# 過去走特徴量（shift/rolling）
for col in ['着順', '人気', '1着差', '上り', '距離']:
    df[f'1走前_{col}'] = df.groupby('馬名')[col].shift(1).fillna(0)

g = df.groupby('馬名')
df['直近5走_平均着順'] = g['着順'].transform(lambda s: s.shift(1).rolling(5, min_periods=1).mean())
df['直近5走_3着内率']  = g['is_top3'].transform(lambda s: s.shift(1).rolling(5, min_periods=1).mean())
# ... (他のrolling特徴量も同様)

# 2019-2024 に絞り split 割り当て
df = df[df['日付'].between(20190101, 20241231)]
df['split'] = 'train'
df.loc[df['日付'].between(20230101, 20231231), 'split'] = 'val'
df.loc[df['日付'].between(20240101, 20241231), 'split'] = 'test'

# レース単位テンソルをベクトル演算で構築
races_df['race_idx'] = np.arange(len(races_df))
df['pos'] = df.groupby('race_idx').cumcount()   # 馬番昇順 position

X_num  = np.zeros((N_RACES, MAX_HORSES, NUM_DIM),  dtype=np.float32)
X_mask = np.zeros((N_RACES, MAX_HORSES),            dtype=bool)
X_num[df['race_idx'], df['pos'], :]  = df[NUM_FEAT_NAMES].values
X_mask[df['race_idx'], df['pos']]    = True

# 勝ち馬 index（同着は最小 pos を採用）
winners = df[df['着順'] == 1].groupby('race_idx')['pos'].min()
win_idx[winners.index] = winners.values

# 02_train_self_attention_model.py（モデル定義の核心部分）

import torch, torch.nn as nn, torch.nn.functional as F

class HorseFeatureEncoder(nn.Module):
    """各馬の数値+カテゴリ特徴量 → d_model 次元"""
    def __init__(self, num_dim, cat_cardinalities, cat_emb_dim=8, d_model=64, dropout=0.1):
        super().__init__()
        self.embeddings = nn.ModuleList([
            nn.Embedding(card, cat_emb_dim, padding_idx=0)
            for card in cat_cardinalities
        ])
        in_dim = num_dim + cat_emb_dim * len(cat_cardinalities)
        self.proj = nn.Sequential(nn.Linear(in_dim, d_model), nn.ReLU(), nn.Dropout(dropout))

    def forward(self, x_num, x_cat):
        parts = [x_num] + [emb(x_cat[:, :, i]) for i, emb in enumerate(self.embeddings)]
        return self.proj(torch.cat(parts, dim=-1))   # (B, H, d_model)


class RaceSelfAttentionModel(nn.Module):
    """
    1レース内の全馬トークンに Self-Attention をかける。
    - positional encoding なし（出走馬の集合として扱う）
    - src_key_padding_mask でパディング馬を除外
    - TransformerEncoder 出力 → 32次元 race_contextual_embedding
    """
    def __init__(self, num_dim, cat_cardinalities,
                 cat_emb_dim=8, d_model=64, nhead=4,
                 num_layers=2, dim_feedforward=128, dropout=0.1, ctx_dim=32):
        super().__init__()
        self.horse_encoder = HorseFeatureEncoder(...)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward,
            dropout=dropout, batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.ctx_proj = nn.Linear(d_model, ctx_dim)
        self.head = nn.Sequential(
            nn.Linear(ctx_dim, 32), nn.ReLU(), nn.Dropout(dropout), nn.Linear(32, 3)
        )

    def forward(self, x_num, x_cat, mask):
        h = self.horse_encoder(x_num, x_cat)                        # (B, H, d_model)
        h = self.transformer(h, src_key_padding_mask=~mask)         # (B, H, d_model)
        ctx_emb    = self.ctx_proj(h)                               # (B, H, ctx_dim=32)
        logits     = self.head(ctx_emb)                             # (B, H, 3)
        win_score  = logits[:, :, 0].masked_fill(~mask, float('-inf'))
        top2_logit = logits[:, :, 1]
        top3_logit = logits[:, :, 2]
        return win_score, top2_logit, top3_logit, ctx_emb

# 損失計算
def compute_loss(win_score, top2_logit, top3_logit, y, win_idx, mask):
    # win: CrossEntropy（softmax はここで内包）
    loss_win = F.cross_entropy(win_score, win_idx)

    # top2/top3: BCE（有効馬のみ）
    flat = mask.view(-1)
    loss_top2 = F.binary_cross_entropy_with_logits(
        top2_logit.view(-1)[flat], y[:,:,1].reshape(-1)[flat], pos_weight=pw_top2
    )
    loss_top3 = F.binary_cross_entropy_with_logits(
        top3_logit.view(-1)[flat], y[:,:,2].reshape(-1)[flat], pos_weight=pw_top3
    )
    return loss_win + loss_top2 + loss_top3