RecSys2023のTurorialを元に頻度主義A/B , ベイズA/B, Bandit A/Bについて理解する

from typing import Dict, List, Union
import numpy as np
import pandas as pd
import itertools

from scipy import stats
from scipy.stats import beta, gamma

from bayesian_test import Bayesian_AB_Test


class Environment:
    def __init__(self, env_params: Dict, config: Dict) -> None:
        self.env_params = env_params
        self.n_impr_per_period = config["max_impr_before_update_param"]

    def simulate_impr_one_period(self) -> pd.DataFrame:
        """Generate random number of impression from environment"""
        n_impr = np.random.randint(0, self.n_impr_per_period)
        return n_impr

    def simulate_clicks_one_period(self, variant: str, n_impr: int) -> pd.DataFrame:
        """Generate random clicks from environment
        Returns:
            n_clicks (int): number of clicks
        """
        n_clicks = np.random.binomial(n_impr, self.env_params[variant]["ctr"])
        return n_clicks

    def sim_clicks_cost_one_period(self, selected_variants: Dict) -> pd.DataFrame:
        """Simulate campaign observation stastistics for one period"""
        variants = list(selected_variants.keys())
        response = {}

        # Simulate impressions and clicks for all variants
        for variant in variants:
            ctr = self.env_params[variant]["ctr"]
            n_impr = selected_variants[variant]
            response.update(
                {
                    variant: {
                        "clicks": np.random.binomial(n_impr, ctr),
                    }
                }
            )

        return response


# =============================================================================
class Agent:

    def __init__(self, variants: List, config: Dict) -> None:
        """Inits class - store params locally
        """
        # Init name of variants
        self.variants = np.array(variants)

        # Init variants with zero impressions and clicks
        self.df_log = {}
        for variant in variants:
            self.create_log(variant, period=0)

        self.optimise_for = config["optimise_for"]
        self.recency_param = config["recency_param"]
        self.n_periods_per_day = config["n_periods_per_day"]

    def create_log(self, variant: str, period: int = 0) -> None:
        """Create log for agent

        Args:
            variant (str): name of the variant
            period (int): current period
        """
        self.df_log[variant] = pd.DataFrame(
            {
                "period": period,
                "n_impr": 0,
                "n_impr_w_sum": 0,
                "n_clicks": 0,
                "n_clicks_w_sum": 0,
                "ctr": 0,
                "alpha": 1,
                "beta": 1,
                "a": 1,
                "scale": 1000,
            },
            index=[0],
        )

    def update_recency(self, variant: str, period: int) -> None:
        """Update recency for each variant, both impressions and clicks
            As current timestanmpo can change, it updates all rows on only the following columns:
                - delta
                - recency
                - n_impr_w
                - n_clicks_w

        Args:
            variant (str): name of the variant
            period (int): current period

        """
        # Calc. recency
        self.df_log[variant]["delta"] = (
            period - self.df_log[variant]["period"]
        ) // self.n_periods_per_day
        self.df_log[variant]["recency"] = (
            self.recency_param ** self.df_log[variant]["delta"]
        )
        self.df_log[variant]["n_impr_w"] = (
            self.df_log[variant]["n_impr"] * self.df_log[variant]["recency"]
        )
        self.df_log[variant]["n_clicks_w"] = (
            self.df_log[variant]["n_clicks"] * self.df_log[variant]["recency"]
        )

    def choose(self, period: int, n_impr: int) -> Union[str, int]:
        """Choose which variants to serve next
        Returns:
            selected_variants (str): names of the variant to serve
        """
        samples = {}

        # 過去の状況において、n_imprしてもらったときに、A,Bどちらが何回勝つかを計算する
        # それは、ベータ分布からサンプリングすることで計算できる
        for variant in self.variants:
            # Get alpha and beta from previous period
            df_prev = self.df_log[variant][self.df_log[variant]["period"] == period - 1]
            # Update recency
            samples[variant] = stats.beta(a=df_prev.alpha, b=df_prev.beta).rvs(n_impr)

        # Thompson sampling
        selected_variants = np.array(
            [samples[variant] for variant in self.variants]
        ).argmax(axis=0)

        selected_variants = {
            variant: sum(self.variants[selected_variants] == variant)
            for variant in self.variants
        }
        print(f"sampled: {samples}")

        return selected_variants

    def update_log(
        self, period: int, selected_variants: Dict, env_response: Dict
    ) -> None:
        """Update log for each variant

        Args:
            period (int): current period
            selected_variants (Dict): Dictionary holding all selected variants
            env_response (Dict): Dictionary holding all environment responses

        Returns:
            log (Dict): Dictionary holding all log entries
        """
        log = {}
        # Update log for each variant
        for variant in self.variants:
            # Calc. 現在のperiodと各peridoの差分を計算
            delta = (period - self.df_log[variant]["period"]) // self.n_periods_per_day
            recency_decay = self.recency_param**delta


            # 過去のデータほど重みを小さくして、和を取る
            n_impr_w_sum = sum(self.df_log[variant]["n_impr"] * recency_decay) + 1
            n_clicks_w_sum = sum(self.df_log[variant]["n_clicks"] * recency_decay) + 1

            ctr = n_clicks_w_sum / n_impr_w_sum
            alpha = np.round(n_clicks_w_sum + 1).astype(int)
            beta = np.round(n_impr_w_sum - n_clicks_w_sum + 1).astype(int)

            scale = 1 / (n_clicks_w_sum + 1e-16)

            # Add newest raw entries to log
            tmp = {
                "period": period,
                "n_impr": selected_variants[variant],  # current impressions
                "n_impr_w_sum": n_impr_w_sum,  # accumulated weighted impressions from all previous periods
                "n_clicks": env_response[variant]["clicks"],
                "n_clicks_w_sum": n_clicks_w_sum,
                "ctr": ctr,
                # For plotting the beta & gamma distributions
                "alpha": alpha,
                "beta": beta,
                "scale": scale,
            }
            # Append as additional row to df_log
            df_tmp = pd.DataFrame(tmp, index=[0])
            self.df_log[variant] = pd.concat(
                [self.df_log[variant], df_tmp], ignore_index=True
            )

            # append to history of agents log
            log[variant] = self.df_log[variant][
                self.df_log[variant]["period"] == period
            ].to_dict(orient="records")[0]

        return log


# =============================================================================
class Bandit:
    """Class running the Bandit in a simulated environment"""

    def __init__(self, bandit_params: Dict, n_periods: int, config: Dict) -> None:
        self.params = bandit_params

        # Extract variants with period 0
        variants = [
            variant for variant in self.params if self.params[variant]["period"] == 0
        ]
        # Set total number of period for the experiment
        self.n_periods = n_periods
        # Extract active variant with highest ctr
        self.optimal_variant = max(
            bandit_params, key=lambda x: bandit_params[x]["ctr"] if x in variants else 0
        )

        # Init environment and agent
        self.environment = Environment(env_params=self.params, config=config)
        self.agent = Agent(variants=variants, config=config)

        # Utility functions
        self.bayes = Bayesian_AB_Test()

    def add_variant(self, variant: str, period: int, env_param: Dict) -> None:
        """Cold start - add another variant suddenly

        Args:
            variant (str): name of the variant
            period (int): current period
            env_param (Dict): Dictionary holding all environment parameters
        """
        # Add new variant to agent and log
        self.agent.variants = np.append(self.agent.variants, variant)
        self.agent.create_log(variant, period=period)
        self.environment.env_params[variant] = env_param
        self.optimal_variant = max(
            self.params,
            key=lambda x: self.params[x]["ctr"] if x in self.agent.variants else 0,
        )

    def run(self):
        """Run the bandit
        Iterate over all periods and update agent and environment
        Calculate metrics
        Inject new variants at specific periods
        """
        self.df_metrics = pd.DataFrame()
        for period in range(1, self.n_periods + 1):
            # ENVIRONMENT - GET IMPRESSIONS
            # random, how many impr. before updating agent
            n_impr = self.environment.simulate_impr_one_period()

            # AGENT - CHOOSE
            # ex) {'A': 35, 'B': 42}
            selected_variants = self.agent.choose(n_impr=n_impr, period=period)

            # ENVIRONMENT - RETURN CLICKS
            # ctrに基づいてimp数を元にclick数を計算する
            response = self.environment.sim_clicks_cost_one_period(selected_variants)
            print(f"period: {period} - selected_variants: {selected_variants}, response: {response}")

            # AGENT - UPDATE LOG
            log = self.agent.update_log(
                period=period,
                selected_variants=selected_variants,
                env_response=response,
            )

            # METRICS
            # Regret
            regret = sum(
                [
                    selected_variants[k]
                    for k in selected_variants
                    if k != self.optimal_variant
                ]
            )

            # Calc. probabilities and losses
            P_ab_ctr = self.bayes.p_ab(
                [
                    beta(log[variant]["alpha"], log[variant]["beta"])
                    for variant in self.agent.variants
                ]
            )

            loss = {"ctr": {}}
            p_overlap = {"ctr": {"ks": {}, "ws": {}}}
            for a, b in itertools.combinations(self.agent.variants, 2):
                loss["ctr"][a, b] = self.bayes.loss(
                    beta(log[a]["alpha"], log[a]["beta"]),
                    beta(log[b]["alpha"], log[b]["beta"]),
                )

            # append df_tmp to df_metrics
            df_tmp = {
                "period": period,
                "n_impr": n_impr,
                "regret": regret,
                "P_ab_ctr": [P_ab_ctr],
                "loss_ctr": [loss["ctr"]],
                "p_overlap_ctr": [p_overlap["ctr"]],
            }
            self.df_metrics = pd.concat(
                [self.df_metrics, pd.DataFrame(df_tmp, index=[0])], ignore_index=True
            )

            self.df_metrics["n_impr_acc"] = self.df_metrics["n_impr"].cumsum()
            self.df_metrics["regret_acc"] = self.df_metrics["regret"].cumsum()
            self.df_metrics["regret_avg"] = (
                self.df_metrics.regret_acc / self.df_metrics.n_impr_acc
            )

            # AGENT - cold start - add another variant suddenly
            for variant in self.params:
                if self.params[variant]["period"] == period:
                    self.add_variant(
                        variant=variant, period=period, env_param=self.params[variant]
                    )