A/B Test Runner in Machine Learning

import hashlib
from typing import Literal


def assign_variant(
    user_id: str,
    experiment_salt: str,
    treatment_fraction: float = 0.5,
) -> Literal["control", "treatment"]:
    """Deterministically assign a user to a variant using hashing.
    
    The hash ensures:
    1. Same user always gets the same variant (consistency)
    2. Assignment is independent across experiments (via salt)
    3. Distribution is uniform (hash output is pseudo-random)
    """
    hash_input = f"{user_id}:{experiment_salt}"
    hash_value = hashlib.md5(hash_input.encode()).hexdigest()
    # Convert first 8 hex chars to a float in [0, 1)
    bucket = int(hash_value[:8], 16) / 0xFFFFFFFF
    return "treatment" if bucket < treatment_fraction else "control"


# Example usage
users = ["user_001", "user_002", "user_003", "user_004", "user_005"]
experiment = "rec_model_v2_2026q1"

for uid in users:
    variant = assign_variant(uid, experiment)
    print(f"{uid} -> {variant}")

# Verify consistency: same user always gets the same variant
assert assign_variant("user_001", experiment) == assign_variant("user_001", experiment)

# Verify independence: different experiments assign differently
result_exp1 = assign_variant("user_001", "experiment_1")
result_exp2 = assign_variant("user_001", "experiment_2")
print(f"\nUser 001 in exp1: {result_exp1}, exp2: {result_exp2}")

import math
from scipy import stats
from typing import Optional


def calculate_sample_size(
    baseline_rate: float,
    mde_relative: float,
    alpha: float = 0.05,
    power: float = 0.80,
    two_sided: bool = True,
) -> dict:
    """Calculate required sample size per group for a proportion test.
    
    Args:
        baseline_rate: Current conversion rate (e.g., 0.05 for 5%)
        mde_relative: Minimum detectable effect as relative lift (e.g., 0.10 for 10%)
        alpha: Significance level (Type I error rate)
        power: Statistical power (1 - Type II error rate)
        two_sided: Whether to use a two-sided test
    
    Returns:
        Dictionary with sample size and experiment parameters
    """
    # Calculate absolute MDE from relative lift
    mde_absolute = baseline_rate * mde_relative
    
    # Pooled variance for proportions
    p_control = baseline_rate
    p_treatment = baseline_rate + mde_absolute
    
    # Z-scores
    z_alpha = stats.norm.ppf(1 - alpha / (2 if two_sided else 1))
    z_beta = stats.norm.ppf(power)
    
    # Variance under null (pooled) and alternative
    var_null = 2 * p_control * (1 - p_control)
    var_alt = (p_control * (1 - p_control)) + (p_treatment * (1 - p_treatment))
    
    # Sample size per group
    n_per_group = ((z_alpha * math.sqrt(var_null) + z_beta * math.sqrt(var_alt)) 
                   / mde_absolute) ** 2
    n_per_group = math.ceil(n_per_group)
    
    return {
        "n_per_group": n_per_group,
        "n_total": n_per_group * 2,
        "baseline_rate": baseline_rate,
        "mde_relative": f"{mde_relative:.1%}",
        "mde_absolute": f"{mde_absolute:.4f}",
        "expected_treatment_rate": p_treatment,
        "alpha": alpha,
        "power": power,
    }


# Example 1: E-commerce conversion rate
result = calculate_sample_size(
    baseline_rate=0.032,     # 3.2% conversion rate
    mde_relative=0.10,       # Detect a 10% relative lift
    alpha=0.05,
    power=0.80,
)
print("E-commerce Conversion Test:")
print(f"  Need {result['n_per_group']:,} users per group")
print(f"  Total: {result['n_total']:,} users")
print(f"  Detecting: {result['baseline_rate']:.1%} -> {result['expected_treatment_rate']:.1%}")

# Example 2: CTR improvement for a recommendation model
result2 = calculate_sample_size(
    baseline_rate=0.12,      # 12% CTR
    mde_relative=0.05,       # Detect a 5% relative lift
    alpha=0.05,
    power=0.80,
)
print(f"\nRecommendation CTR Test:")
print(f"  Need {result2['n_per_group']:,} users per group")
print(f"  Total: {result2['n_total']:,} users")

import numpy as np
from scipy import stats
from dataclasses import dataclass
from typing import Optional


@dataclass
class ABTestResult:
    control_mean: float
    treatment_mean: float
    absolute_lift: float
    relative_lift: float
    p_value: float
    confidence_interval: tuple
    is_significant: bool
    n_control: int
    n_treatment: int
    variance_reduction: Optional[float] = None


def analyze_ab_test(
    control: np.ndarray,
    treatment: np.ndarray,
    alpha: float = 0.05,
    pre_experiment_control: Optional[np.ndarray] = None,
    pre_experiment_treatment: Optional[np.ndarray] = None,
) -> ABTestResult:
    """Analyze an A/B test with optional CUPED variance reduction.
    
    Args:
        control: Metric values for the control group
        treatment: Metric values for the treatment group
        alpha: Significance level
        pre_experiment_control: Pre-experiment covariate for control (for CUPED)
        pre_experiment_treatment: Pre-experiment covariate for treatment (for CUPED)
    
    Returns:
        ABTestResult with statistical analysis
    """
    variance_reduction = None
    
    # Apply CUPED if pre-experiment data is provided
    if pre_experiment_control is not None and pre_experiment_treatment is not None:
        # Combine for theta estimation
        all_y = np.concatenate([control, treatment])
        all_x = np.concatenate([pre_experiment_control, pre_experiment_treatment])
        
        # Compute theta = Cov(Y, X) / Var(X)
        theta = np.cov(all_y, all_x)[0, 1] / np.var(all_x)
        
        # Adjust outcomes
        x_mean = np.mean(all_x)
        control_adj = control - theta * (pre_experiment_control - x_mean)
        treatment_adj = treatment - theta * (pre_experiment_treatment - x_mean)
        
        # Calculate variance reduction
        original_var = np.var(np.concatenate([control, treatment]))
        adjusted_var = np.var(np.concatenate([control_adj, treatment_adj]))
        variance_reduction = 1 - (adjusted_var / original_var)
        
        control = control_adj
        treatment = treatment_adj
    
    # Compute statistics
    n_c, n_t = len(control), len(treatment)
    mean_c, mean_t = np.mean(control), np.mean(treatment)
    var_c, var_t = np.var(control, ddof=1), np.var(treatment, ddof=1)
    
    # Standard error of the difference
    se = np.sqrt(var_c / n_c + var_t / n_t)
    
    # Z-statistic and p-value (two-sided)
    z_stat = (mean_t - mean_c) / se
    p_value = 2 * (1 - stats.norm.cdf(abs(z_stat)))
    
    # Confidence interval
    z_crit = stats.norm.ppf(1 - alpha / 2)
    ci_lower = (mean_t - mean_c) - z_crit * se
    ci_upper = (mean_t - mean_c) + z_crit * se
    
    return ABTestResult(
        control_mean=mean_c,
        treatment_mean=mean_t,
        absolute_lift=mean_t - mean_c,
        relative_lift=(mean_t - mean_c) / mean_c if mean_c != 0 else float('inf'),
        p_value=p_value,
        confidence_interval=(ci_lower, ci_upper),
        is_significant=p_value < alpha,
        n_control=n_c,
        n_treatment=n_t,
        variance_reduction=variance_reduction,
    )


# Simulate an A/B test for a recommendation model
np.random.seed(42)
n = 10000

# Pre-experiment behavior (e.g., last week's CTR per user)
pre_control = np.random.beta(2, 18, n)  # ~10% baseline
pre_treatment = np.random.beta(2, 18, n)

# Experiment outcomes (treatment has a 5% relative lift)
control_outcomes = pre_control + np.random.normal(0, 0.05, n)
treatment_outcomes = pre_treatment * 1.05 + np.random.normal(0, 0.05, n)

# Analyze without CUPED
result_raw = analyze_ab_test(control_outcomes, treatment_outcomes)
print("Without CUPED:")
print(f"  Lift: {result_raw.relative_lift:.2%}")
print(f"  P-value: {result_raw.p_value:.4f}")
print(f"  Significant: {result_raw.is_significant}")

# Analyze with CUPED
result_cuped = analyze_ab_test(
    control_outcomes, treatment_outcomes,
    pre_experiment_control=pre_control,
    pre_experiment_treatment=pre_treatment,
)
print(f"\nWith CUPED:")
print(f"  Lift: {result_cuped.relative_lift:.2%}")
print(f"  P-value: {result_cuped.p_value:.6f}")
print(f"  Significant: {result_cuped.is_significant}")
print(f"  Variance Reduction: {result_cuped.variance_reduction:.1%}")

import numpy as np
from scipy import stats


def bayesian_ab_test(
    control_successes: int,
    control_trials: int,
    treatment_successes: int,
    treatment_trials: int,
    prior_alpha: float = 1.0,
    prior_beta: float = 1.0,
    n_simulations: int = 100000,
) -> dict:
    """Bayesian A/B test for binary outcomes using Beta-Binomial model.
    
    Uses conjugate Beta prior: Beta(alpha, beta) with Binomial likelihood
    gives Beta(alpha + successes, beta + failures) posterior.
    
    Args:
        control_successes: Number of conversions in control
        control_trials: Total users in control
        treatment_successes: Number of conversions in treatment
        treatment_trials: Total users in treatment
        prior_alpha: Beta prior alpha (1.0 = uniform/non-informative)
        prior_beta: Beta prior beta (1.0 = uniform/non-informative)
        n_simulations: Number of Monte Carlo samples
    
    Returns:
        Dictionary with posterior statistics and decision metrics
    """
    # Posterior parameters (conjugate update)
    control_a = prior_alpha + control_successes
    control_b = prior_beta + (control_trials - control_successes)
    treatment_a = prior_alpha + treatment_successes
    treatment_b = prior_beta + (treatment_trials - treatment_successes)
    
    # Sample from posteriors
    control_samples = np.random.beta(control_a, control_b, n_simulations)
    treatment_samples = np.random.beta(treatment_a, treatment_b, n_simulations)
    
    # Probability that treatment is better
    prob_treatment_better = np.mean(treatment_samples > control_samples)
    
    # Expected lift distribution
    lift_samples = (treatment_samples - control_samples) / control_samples
    
    # Risk: expected loss if we choose treatment but it's actually worse
    loss_if_choose_treatment = np.mean(
        np.maximum(control_samples - treatment_samples, 0)
    )
    loss_if_choose_control = np.mean(
        np.maximum(treatment_samples - control_samples, 0)
    )
    
    return {
        "prob_treatment_better": prob_treatment_better,
        "expected_lift": np.mean(lift_samples),
        "lift_95_ci": (np.percentile(lift_samples, 2.5), np.percentile(lift_samples, 97.5)),
        "control_posterior_mean": control_a / (control_a + control_b),
        "treatment_posterior_mean": treatment_a / (treatment_a + treatment_b),
        "risk_choosing_treatment": loss_if_choose_treatment,
        "risk_choosing_control": loss_if_choose_control,
        "recommended_action": "Ship treatment" if prob_treatment_better > 0.95 else 
                               "Ship control" if prob_treatment_better < 0.05 else
                               "Continue experiment",
    }


# Example: Testing a new ML-based search ranking at an Indian e-commerce platform
np.random.seed(42)

result = bayesian_ab_test(
    control_successes=3200,    # 3200 conversions out of 100k
    control_trials=100000,
    treatment_successes=3350,  # 3350 conversions out of 100k
    treatment_trials=100000,
)

print("Bayesian A/B Test Results:")
print(f"  Control conversion rate:   {result['control_posterior_mean']:.3%}")
print(f"  Treatment conversion rate: {result['treatment_posterior_mean']:.3%}")
print(f"  P(treatment > control):    {result['prob_treatment_better']:.1%}")
print(f"  Expected relative lift:    {result['expected_lift']:.2%}")
print(f"  95% CI for lift:           [{result['lift_95_ci'][0]:.2%}, {result['lift_95_ci'][1]:.2%}]")
print(f"  Risk if ship treatment:    {result['risk_choosing_treatment']:.5f}")
print(f"  Risk if ship control:      {result['risk_choosing_control']:.5f}")
print(f"  Recommendation:            {result['recommended_action']}")

import numpy as np
from typing import List
from dataclasses import dataclass, field


@dataclass
class BanditArm:
    name: str
    alpha: float = 1.0  # Beta prior successes
    beta: float = 1.0   # Beta prior failures
    total_pulls: int = 0
    total_rewards: int = 0
    
    def update(self, reward: int):
        self.alpha += reward
        self.beta += (1 - reward)
        self.total_pulls += 1
        self.total_rewards += reward
    
    def sample(self) -> float:
        return np.random.beta(self.alpha, self.beta)
    
    @property
    def conversion_rate(self) -> float:
        if self.total_pulls == 0:
            return 0.0
        return self.total_rewards / self.total_pulls


def run_thompson_sampling(
    true_rates: List[float],
    arm_names: List[str],
    n_rounds: int = 10000,
) -> dict:
    """Run Thompson Sampling multi-armed bandit.
    
    Thompson Sampling is a Bayesian approach that balances
    exploration (trying uncertain arms) and exploitation
    (choosing the best-known arm) naturally through posterior
    sampling.
    """
    arms = [BanditArm(name=name) for name in arm_names]
    cumulative_regret = []
    total_regret = 0
    best_rate = max(true_rates)
    
    for round_num in range(n_rounds):
        # Sample from each arm's posterior and pick the highest
        samples = [arm.sample() for arm in arms]
        chosen_idx = np.argmax(samples)
        
        # Simulate reward from the true (unknown) conversion rate
        reward = np.random.binomial(1, true_rates[chosen_idx])
        arms[chosen_idx].update(reward)
        
        # Track regret (difference from always choosing the best arm)
        total_regret += best_rate - true_rates[chosen_idx]
        cumulative_regret.append(total_regret)
    
    return {
        "arms": [
            {
                "name": arm.name,
                "pulls": arm.total_pulls,
                "observed_rate": f"{arm.conversion_rate:.3%}",
                "traffic_share": f"{arm.total_pulls / n_rounds:.1%}",
            }
            for arm in arms
        ],
        "total_regret": total_regret,
        "avg_regret_per_round": total_regret / n_rounds,
    }


# Example: Testing 3 recommendation models simultaneously
np.random.seed(42)
result = run_thompson_sampling(
    true_rates=[0.10, 0.12, 0.11],  # True CTRs (unknown to the algorithm)
    arm_names=["Model A (baseline)", "Model B (new)", "Model C (experimental)"],
    n_rounds=10000,
)

print("Thompson Sampling Results (10,000 rounds):")
print(f"{'Arm':<30} {'Pulls':>8} {'Rate':>10} {'Traffic':>10}")
print("-" * 60)
for arm in result["arms"]:
    print(f"{arm['name']:<30} {arm['pulls']:>8} {arm['observed_rate']:>10} {arm['traffic_share']:>10}")
print(f"\nTotal regret: {result['total_regret']:.1f}")
print(f"Avg regret/round: {result['avg_regret_per_round']:.4f}")

# GrowthBook experiment configuration (YAML)
experiment:
  id: rec-model-v2-2026q1
  name: Recommendation Model V2 Test
  hypothesis: >-
    The new collaborative filtering model (V2) will increase
    add-to-cart rate by at least 5% relative to the current
    content-based model (V1).
  
  # Traffic allocation
  allocation:
    control: 0.50     # 50% traffic
    treatment: 0.50   # 50% traffic
  
  # Targeting
  targeting:
    platforms: [web, android, ios]
    countries: [IN, US, UK]
    user_segments: [returning_users]
  
  # Metrics
  metrics:
    primary:
      - name: add_to_cart_rate
        type: proportion
        direction: increase
    secondary:
      - name: revenue_per_user
        type: mean
        direction: increase
      - name: items_viewed
        type: mean
        direction: increase
    guardrails:
      - name: p99_latency_ms
        threshold: 500
        direction: must_not_increase
      - name: error_rate
        threshold: 0.01
        direction: must_not_increase
      - name: crash_rate
        threshold: 0.001
        direction: must_not_increase
  
  # Analysis settings
  analysis:
    method: frequentist_sequential  # or: bayesian, frequentist_fixed
    alpha: 0.05
    power: 0.80
    mde_relative: 0.05              # 5% relative lift
    variance_reduction: cuped
    cuped_covariate: pre_experiment_add_to_cart_rate
    multiple_comparison_correction: benjamini_hochberg
  
  # Duration
  schedule:
    start_date: 2026-03-01
    min_runtime_days: 14
    max_runtime_days: 42
    estimated_sample_size_per_group: 50000