get_kl.py

#!/usr/bin/env python3
import argparse
import os
from collections import Counter, defaultdict
from typing import Dict, Iterable, List, Sequence, Tuple

import numpy as np
from scipy.special import rel_entr


def float_to_str(f: float) -> str:
    return f"{int(f * 10):02d}"


def load_sequences(path: str) -> Tuple[List[List[int]], List[List[int]]]:
    with open(path, "rb") as f:
        train, test = pickle_load(f)
    # Ensure python lists (not np arrays) for fast iteration
    train = [list(seq) for seq in train]
    test = [list(seq) for seq in test]
    return train, test


def pickle_load(f):
    # Inline function so we don't import pickle at module import time
    import pickle as _pkl

    return _pkl.load(f)


def seq_len_ok(seq: Sequence[int], min_len: int = 2) -> bool:
    return len(seq) >= min_len


def build_item_universe(
    train: Iterable[Sequence[int]],
    test: Iterable[Sequence[int]],
    universe: str = "test",
) -> List[int]:
    """
    universe = 'test' (default): only items seen in test (matches original behavior)
             = 'union'        : union of train and test items
    """
    items = set()
    if universe in ("test", "union"):
        for s in test:
            items.update(s)
    if universe == "union":
        for s in train:
            items.update(s)
    return list(items)


def compute_validation_distribution(
    test_sequences: Iterable[Sequence[int]],
    item_list: List[int],
    validation: bool,
    eps: float,
) -> np.ndarray:
    """
    If validation=True, follow your original behavior: drop the last element of every
    test sequence, then count the new last element as the target.
    """
    idx = {item: i for i, item in enumerate(item_list)}
    counts = np.zeros(len(item_list), dtype=np.float64)

    for seq in test_sequences:
        if not seq_len_ok(seq, 1):
            continue
        # If validation, use seq[:-1] as the truncated sequence if possible
        target_seq = seq[:-1] if validation and len(seq) >= 2 else seq
        if not target_seq:
            continue
        y = target_seq[-1]
        j = idx.get(y, None)
        if j is not None:
            counts[j] += 1.0

    total = counts.sum()
    if total == 0:
        # Avoid divide-by-zero; return uniform smoothed distribution over item_list
        probs = np.full_like(counts, 1.0 / len(item_list))
    else:
        probs = counts / total

    probs = probs + eps
    probs /= probs.sum()
    return probs


def compute_alpha_beta_distribution(
    train_sequences: Iterable[Sequence[int]],
    item_list: List[int],
    alpha: float,
    beta: float,
    eps: float,
) -> np.ndarray:
    """
    For each train sequence s = [x1, x2, ..., xL]:
      - sequence weight = (L-1)^alpha
      - positional weights:
          if beta > 10 -> last-target only
          else pos_weights[i] = i^beta for i in 1..L-1 (with pos_weights[0]=0)
      - each target at position i contributes seq_weight * (pos_weight_i / sum_pos_weights)
    """
    idx = {item: i for i, item in enumerate(item_list)}
    contrib = np.zeros(len(item_list), dtype=np.float64)

    for seq in train_sequences:
        L = len(seq)
        if L <= 1:
            continue

        seq_weight = (L - 1) ** alpha

        if beta > 10.0:
            # only last position
            pos_weights = np.zeros(L, dtype=np.float64)
            pos_weights[-1] = 1.0
        else:
            pos_weights = np.zeros(L, dtype=np.float64)
            # indices 1..L-1 are target positions in your original code
            # i^beta with i starting at 1 for the second element
            pos_weights[1:] = np.power(np.arange(1, L, dtype=np.float64), beta)

        wsum = pos_weights.sum()
        if wsum <= 0:
            continue

        # Distribute weight to item at each target position
        for i in range(1, L):
            w = pos_weights[i] / wsum
            y = seq[i]
            j = idx.get(y, None)
            if j is not None:
                contrib[j] += seq_weight * w

    total = contrib.sum()
    if total == 0:
        probs = np.full_like(contrib, 1.0 / len(item_list))
    else:
        probs = contrib / total

    probs = probs + eps
    probs /= probs.sum()
    return probs


def kl_divergence(p: np.ndarray, q: np.ndarray) -> float:
    """
    KL(p || q) = sum p * log(p/q); numpy-safe via scipy.special.rel_entr
    Assumes p and q are strictly positive and normalized.
    """
    return float(np.sum(rel_entr(p, q)))


def main():
    parser = argparse.ArgumentParser(description="Compute KL between validation and alpha/beta-weighted train target distributions.")
    parser.add_argument("--dataset", type=str, default="beauty")
    parser.add_argument("--alpha", type=float, default=1.0)
    parser.add_argument("--beta", type=float, default=0.0)
    parser.add_argument("--validation", type=lambda x: x.lower() == "true", default=False)
    parser.add_argument("--data_dir", type=str, default="data_pkl", help="Directory containing <dataset>_seq.pkl")
    parser.add_argument("--smoothing_eps", type=float, default=1e-10, help="Additive smoothing before renormalization.")
    parser.add_argument("--item_universe", type=str, choices=["test", "union"], default="test",
                        help="'test' uses only items from test (original behavior), 'union' uses train∪test.")
    parser.add_argument("--seed", type=int, default=42, help="Random seed (for reproducibility if you add randomness later).")
    args = parser.parse_args()

    # Fix random seed (future-proofing if randomness is added)
    import random as _random

    _random.seed(args.seed)
    np.random.seed(args.seed)

    # Load
    path = os.path.join(args.data_dir, f"{args.dataset}_seq.pkl")
    train_sequences, test_sequences = load_sequences(path)

    # Build item universe
    item_list = build_item_universe(train_sequences, test_sequences, universe=args.item_universe)

    # Compute distributions
    p_valid = compute_validation_distribution(test_sequences, item_list, args.validation, args.smoothing_eps)
    q_train = compute_alpha_beta_distribution(train_sequences, item_list, args.alpha, args.beta, args.smoothing_eps)

    # KL
    kl = kl_divergence(p_valid, q_train)

    data_type = 'Validation' if args.validation else 'Test'

    print("\n=== KL Divergence Report ===")
    print(f" Dataset       : {args.dataset}")
    print(f" Alpha (α)     : {args.alpha:.2f}")
    print(f" Beta  (β)     : {args.beta:.2f}")
    print(f" Data Type     : {data_type}")
    print(f" KL(p_valid || q_train) = {kl:.6f}")
    print("============================\n")



if __name__ == "__main__":
    main()