A/B Testing for Machine Learning Models: ML Experimentation

from dataclasses import dataclass, field
from collections import defaultdict
import hashlib
from typing import Any

@dataclass
class MLExperiment:
    name: str
    hypothesis: str
    success_criteria: dict
    oec: str  # Overall Evaluation Criterion
    guardrails: dict = field(default_factory=dict)
    variants: dict = field(default_factory=dict)

    def add_variant(self, name: str, model: Any, traffic_allocation: float):
        self.variants[name] = {
            'model': model,
            'traffic_allocation': traffic_allocation,
            'metrics': defaultdict(list)
        }

# Example
experiment = MLExperiment(
    name="ranking_model_v2",
    hypothesis="Transformer model will increase CTR by 5% without increasing latency",
    oec="click_through_rate",
    success_criteria={
        'ctr_increase': 0.05,
        'latency_p95_max': 200,  # ms
        'min_statistical_power': 0.80,
        'significance_level': 0.05
    },
    guardrails={
        'error_rate_max_increase': 0.10,  # Max 10% increase
        'latency_p99_max': 500,  # Hard SLA limit
        'revenue_max_decrease': 0.02  # Max 2% revenue drop
    }
)

experiment.add_variant('control', model_v1, traffic_allocation=0.5)
experiment.add_variant('treatment', model_v2, traffic_allocation=0.5)

class TrafficSplitter:
    def __init__(self, experiment: MLExperiment):
        self.experiment = experiment

    def assign_variant(self, user_id: str) -> tuple[str, Any]:
        """Consistent hash-based assignment ensures same user always gets same variant"""
        hash_value = hashlib.md5(
            f"{self.experiment.name}:{user_id}".encode()
        ).hexdigest()

        hash_int = int(hash_value, 16)
        threshold = hash_int % 100

        cumulative = 0
        for variant_name, variant in self.experiment.variants.items():
            cumulative += variant['traffic_allocation'] * 100
            if threshold < cumulative:
                return variant_name, variant['model']

        return 'control', self.experiment.variants['control']['model']

# Usage
splitter = TrafficSplitter(experiment)
variant, model = splitter.assign_variant(user_id="user_12345")
prediction = model.predict(features)

from scipy import stats
import numpy as np

class SignificanceTester:
    def __init__(self, alpha: float = 0.05, power: float = 0.80):
        self.alpha = alpha
        self.power = power

    def calculate_sample_size(self, baseline_rate: float, mde: float) -> int:
        """Calculate minimum sample size per variant for given MDE"""
        # Using formula for two-proportion z-test
        from scipy.stats import norm

        z_alpha = norm.ppf(1 - self.alpha / 2)
        z_beta = norm.ppf(self.power)

        p1 = baseline_rate
        p2 = baseline_rate * (1 + mde)
        p_avg = (p1 + p2) / 2

        n = (2 * p_avg * (1 - p_avg) * (z_alpha + z_beta) ** 2) / ((p2 - p1) ** 2)
        return int(np.ceil(n))

    def t_test(self, control_data: list, treatment_data: list) -> dict:
        """Two-sample t-test for continuous metrics"""
        t_stat, p_value = stats.ttest_ind(control_data, treatment_data)

        return {
            't_statistic': t_stat,
            'p_value': p_value,
            'is_significant': p_value < self.alpha,
            'control_mean': np.mean(control_data),
            'treatment_mean': np.mean(treatment_data),
            'relative_lift': (np.mean(treatment_data) - np.mean(control_data)) / np.mean(control_data),
            'confidence_interval': self._confidence_interval(control_data, treatment_data)
        }

    def chi_square_test(self, control_conversions: int, control_total: int,
                        treatment_conversions: int, treatment_total: int) -> dict:
        """Chi-square test for binary metrics (clicks, conversions)"""
        contingency_table = np.array([
            [control_conversions, control_total - control_conversions],
            [treatment_conversions, treatment_total - treatment_conversions]
        ])

        chi2, p_value, dof, expected = stats.chi2_contingency(contingency_table)

        control_rate = control_conversions / control_total
        treatment_rate = treatment_conversions / treatment_total

        return {
            'chi2_statistic': chi2,
            'p_value': p_value,
            'is_significant': p_value < self.alpha,
            'control_rate': control_rate,
            'treatment_rate': treatment_rate,
            'relative_lift': (treatment_rate - control_rate) / control_rate
        }

    def _confidence_interval(self, control: list, treatment: list, confidence: float = 0.95):
        """Calculate confidence interval for the difference"""
        diff = np.mean(treatment) - np.mean(control)
        se = np.sqrt(np.var(control)/len(control) + np.var(treatment)/len(treatment))
        z = stats.norm.ppf((1 + confidence) / 2)
        return (diff - z * se, diff + z * se)

# Example usage
tester = SignificanceTester(alpha=0.05)

# Calculate required sample size before running experiment
sample_size = tester.calculate_sample_size(
    baseline_rate=0.12,  # 12% current CTR
    mde=0.05  # Want to detect 5% improvement
)
print(f"Need {sample_size:,} samples per variant")

# After collecting data
ctr_result = tester.chi_square_test(
    control_conversions=1250,
    control_total=10000,
    treatment_conversions=1400,
    treatment_total=10000
)
print(f"CTR lift: {ctr_result['relative_lift']:.2%}")
print(f"Statistically significant: {ctr_result['is_significant']}")

from sklearn.metrics import roc_auc_score

class OfflineEvaluator:
    def __init__(self, test_data: dict):
        self.X_test = test_data['X_test']
        self.y_test = test_data['y_test']

    def holdout_validation(self, model_old, model_new) -> dict:
        """Compare models on held-out data"""
        old_predictions = model_old.predict_proba(self.X_test)[:, 1]
        new_predictions = model_new.predict_proba(self.X_test)[:, 1]

        old_auc = roc_auc_score(self.y_test, old_predictions)
        new_auc = roc_auc_score(self.y_test, new_predictions)

        return {
            'old_model_auc': old_auc,
            'new_model_auc': new_auc,
            'auc_improvement': new_auc - old_auc,
            'recommendation': 'proceed_to_online' if new_auc > old_auc else 'iterate'
        }

    def replay_evaluation(self, model, logged_data: dict) -> dict:
        """Estimate counterfactual performance using inverse propensity scoring"""
        propensity_scores = logged_data['propensity_scores']
        rewards = logged_data['rewards']
        actions = logged_data['actions']

        new_actions = model.predict(logged_data['features'])

        # Inverse propensity scoring for unbiased estimate
        ips_estimate = np.mean([
            rewards[i] / propensity_scores[i] if new_actions[i] == actions[i] else 0
            for i in range(len(rewards))
        ])

        return {'estimated_reward': ips_estimate}

class OnlineEvaluator:
    def __init__(self, experiment: MLExperiment):
        self.experiment = experiment
        self.alerts = []

    def check_guardrails(self) -> dict:
        """Automatically halt experiment if critical metrics degrade"""
        control = self.experiment.variants['control']['metrics']
        treatment = self.experiment.variants['treatment']['metrics']

        violations = []

        # Check error rate
        control_errors = np.mean(control['errors']) if control['errors'] else 0
        treatment_errors = np.mean(treatment['errors']) if treatment['errors'] else 0

        if control_errors > 0:
            error_increase = (treatment_errors - control_errors) / control_errors
            max_allowed = self.experiment.guardrails.get('error_rate_max_increase', 0.10)
            if error_increase > max_allowed:
                violations.append(f"Error rate +{error_increase:.1%} exceeds {max_allowed:.0%} limit")

        # Check latency
        if treatment['latency_ms']:
            p99_latency = np.percentile(treatment['latency_ms'], 99)
            max_p99 = self.experiment.guardrails.get('latency_p99_max', 500)
            if p99_latency > max_p99:
                violations.append(f"P99 latency {p99_latency:.0f}ms exceeds {max_p99}ms SLA")

        return {
            'passed': len(violations) == 0,
            'violations': violations,
            'action': 'continue' if len(violations) == 0 else 'halt_experiment'
        }

    def real_time_monitoring(self) -> dict:
        """Generate alerts for concerning trends"""
        summary = self._calculate_current_performance()

        alerts = []

        if summary['treatment']['error_rate'] > summary['control']['error_rate'] * 1.5:
            alerts.append({
                'severity': 'CRITICAL',
                'message': f"Treatment error rate 50% higher than control",
                'action': 'Consider halting experiment'
            })

        return {'summary': summary, 'alerts': alerts}

class ThompsonSampling:
    """Adaptive traffic allocation based on observed performance"""

    def __init__(self, variants: list[str]):
        # Beta distribution parameters (prior: uniform)
        self.variants = {
            name: {'alpha': 1, 'beta': 1}
            for name in variants
        }

    def select_variant(self) -> str:
        """Sample from posterior distributions, pick highest"""
        samples = {
            name: np.random.beta(params['alpha'], params['beta'])
            for name, params in self.variants.items()
        }
        return max(samples, key=samples.get)

    def update(self, variant_name: str, reward: float):
        """Update posterior based on observed reward"""
        if reward > 0:
            self.variants[variant_name]['alpha'] += 1
        else:
            self.variants[variant_name]['beta'] += 1

    def get_allocation_probabilities(self) -> dict:
        """Current probability of selecting each variant"""
        total_samples = 10000
        selections = [self.select_variant() for _ in range(total_samples)]
        return {name: selections.count(name) / total_samples for name in self.variants}

# Usage
bandit = ThompsonSampling(['model_v1', 'model_v2', 'model_v3'])

for user in users:
    selected_variant = bandit.select_variant()
    prediction = models[selected_variant].predict(user.features)
    reward = user.interact(prediction)  # 1 if click, 0 otherwise
    bandit.update(selected_variant, reward)

class InterleavedTest:
    """Present results from both models, track which users prefer"""

    def __init__(self, model_a, model_b):
        self.model_a = model_a
        self.model_b = model_b
        self.wins_a = 0
        self.wins_b = 0
        self.ties = 0

    def team_draft_interleaving(self, query, k: int = 10) -> list:
        """Create interleaved result list using team-draft method"""
        results_a = self.model_a.rank(query, top_k=k*2)
        results_b = self.model_b.rank(query, top_k=k*2)

        interleaved = []
        used = set()
        ptr_a, ptr_b = 0, 0

        for i in range(k):
            if i % 2 == 0:  # A's turn
                while ptr_a < len(results_a) and results_a[ptr_a] in used:
                    ptr_a += 1
                if ptr_a < len(results_a):
                    interleaved.append({'item': results_a[ptr_a], 'source': 'A'})
                    used.add(results_a[ptr_a])
                    ptr_a += 1
            else:  # B's turn
                while ptr_b < len(results_b) and results_b[ptr_b] in used:
                    ptr_b += 1
                if ptr_b < len(results_b):
                    interleaved.append({'item': results_b[ptr_b], 'source': 'B'})
                    used.add(results_b[ptr_b])
                    ptr_b += 1

        return interleaved

    def evaluate_clicks(self, interleaved_results: list, clicked_indices: list) -> str:
        """Determine winner based on which model's items got clicked"""
        clicks_a = sum(1 for i in clicked_indices if interleaved_results[i]['source'] == 'A')
        clicks_b = sum(1 for i in clicked_indices if interleaved_results[i]['source'] == 'B')

        if clicks_a > clicks_b:
            self.wins_a += 1
            return 'A'
        elif clicks_b > clicks_a:
            self.wins_b += 1
            return 'B'
        else:
            self.ties += 1
            return 'TIE'

class GradualRollout:
    """Safe, incremental model deployment"""

    STAGES = [
        {'allocation': 0.05, 'min_hours': 24, 'name': 'canary'},
        {'allocation': 0.20, 'min_hours': 48, 'name': 'early_adopters'},
        {'allocation': 0.50, 'min_hours': 48, 'name': 'half_traffic'},
        {'allocation': 1.00, 'min_hours': 0, 'name': 'full_rollout'}
    ]

    def __init__(self, experiment: MLExperiment):
        self.experiment = experiment
        self.current_stage = 0

    def should_advance(self, hours_running: float, metrics: dict) -> dict:
        """Determine if safe to increase traffic"""
        stage = self.STAGES[self.current_stage]

        if hours_running < stage['min_hours']:
            return {
                'advance': False,
                'reason': f"Need {stage['min_hours'] - hours_running:.0f} more hours at {stage['name']} stage"
            }

        if not metrics.get('guardrails_passing', False):
            return {
                'advance': False,
                'reason': 'Guardrails failing—investigate before proceeding'
            }

        # Check for statistical significance if we have enough data
        if metrics.get('is_significant') and metrics.get('positive_lift'):
            return {
                'advance': True,
                'next_stage': self.STAGES[self.current_stage + 1]['name'] if self.current_stage + 1 < len(self.STAGES) else 'complete',
                'next_allocation': self.STAGES[self.current_stage + 1]['allocation'] if self.current_stage + 1 < len(self.STAGES) else 1.0
            }

        return {
            'advance': False,
            'reason': 'Waiting for statistical significance'
        }