federated/aggregator.py

"""Weight aggregation algorithms for Federated Learning.

Implements various aggregation methods beyond simple FedAvg:
- FedProx (proximal term)
- SCAFFOLD (variance reduction)
- FedOpt (adaptive optimization)
"""

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch


class Aggregator(ABC):
    """Base class for federated aggregation algorithms."""
    
    @abstractmethod
    def aggregate(
        self,
        parameters: List[List[np.ndarray]],
        weights: List[float],
        client_metrics: List[Dict],
    ) -> List[np.ndarray]:
        """Aggregate client parameters.
        
        Args:
            parameters: List of client parameters
            weights: Weight for each client (usually proportional to data size)
            client_metrics: Metrics from each client
        
        Returns:
            Aggregated parameters
        """
        pass


class FedAvgAggregator(Aggregator):
    """Federated Averaging (FedAvg) - standard weighted average."""
    
    def aggregate(
        self,
        parameters: List[List[np.ndarray]],
        weights: List[float],
        client_metrics: List[Dict],
    ) -> List[np.ndarray]:
        """Weighted average of client parameters."""
        if not parameters:
            raise ValueError("No parameters to aggregate")
        
        if len(parameters) == 1:
            return parameters[0]
        
        # Normalize weights
        total_weight = sum(weights)
        normalized_weights = [w / total_weight for w in weights]
        
        # Weighted average
        aggregated = None
        
        for params, weight in zip(parameters, normalized_weights):
            if aggregated is None:
                aggregated = [p * weight for p in params]
            else:
                aggregated = [
                    a + p * weight
                    for a, p in zip(aggregated, params)
                ]
        
        return aggregated


class FedProxAggregator(Aggregator):
    """FedProx with proximal term for handling heterogeneity."""
    
    def __init__(self, mu: float = 0.01):
        """
        Args:
            mu: Proximal term coefficient
        """
        self.mu = mu
    
    def aggregate(
        self,
        parameters: List[List[np.ndarray]],
        weights: List[float],
        client_metrics: List[Dict],
        global_parameters: Optional[List[np.ndarray]] = None,
    ) -> List[np.ndarray]:
        """Aggregate with proximal term regularization."""
        if global_parameters is None:
            # Fall back to FedAvg if no global params
            return FedAvgAggregator().aggregate(parameters, weights, client_metrics)
        
        # Weighted average with proximal correction
        total_weight = sum(weights)
        normalized_weights = [w / total_weight for w in weights]
        
        aggregated = None
        
        for params, weight, global_params in zip(
            parameters, normalized_weights, global_parameters
        ):
            # Proximal term: add regularization toward global model
            corrected_params = [
                p + self.mu * (g - p)
                for p, g in zip(params, global_params)
            ]
            
            if aggregated is None:
                aggregated = [p * weight for p in corrected_params]
            else:
                aggregated = [
                    a + p * weight
                    for a, p in zip(aggregated, corrected_params)
                ]
        
        return aggregated


class SCAFFOLDAggregator(Aggregator):
    """SCAFFOLD: Stochastic Controlled Averaging for Federated Learning."""
    
    def __init__(self, global_control: Optional[List[np.ndarray]] = None):
        self.global_control = global_control or None
    
    def set_global_control(self, control: List[np.ndarray]) -> None:
        """Set global control variates."""
        self.global_control = control
    
    def aggregate(
        self,
        parameters: List[List[np.ndarray]],
        weights: List[float],
        client_metrics: List[Dict],
        client_controls: Optional[List[List[np.ndarray]]] = None,
    ) -> List[np.ndarray]:
        """Aggregate using SCAFFOLD algorithm."""
        if client_controls is None or self.global_control is None:
            return FedAvgAggregator().aggregate(parameters, weights, client_metrics)
        
        total_weight = sum(weights)
        normalized_weights = [w / total_weight for w in weights]
        
        # Compute weight updates (gradient-like terms)
        aggregated = None
        
        for params, weight, client_ctrl, global_ctrl in zip(
            parameters, normalized_weights, client_controls, self.global_control
        ):
            # Direction: client_params - global_params + global_control - client_control
            delta = [
                p - g + gc - cc
                for p, g, gc, cc in zip(params, parameters[0], self.global_control, client_ctrl)
            ]
            
            if aggregated is None:
                aggregated = [d * weight for d in delta]
            else:
                aggregated = [
                    a + d * weight
                    for a, d in zip(aggregated, delta)
                ]
        
        # Add back global parameters
        if aggregated:
            aggregated = [
                g + self.mu * a if hasattr(self, 'mu') else g + 0.001 * a
                for g, a in zip(self.global_control if self.global_control else parameters[0], aggregated)
            ]
        
        return aggregated
    
    @property
    def mu(self) -> float:
        """Learning rate for SCAFFOLD."""
        return 0.001


class FedOptAggregator(Aggregator):
    """FedOpt: Adaptive Federated Optimization using server-side optimizer."""
    
    def __init__(
        self,
        server_lr: float = 1.0,
        beta_1: float = 0.9,
        beta_2: float = 0.99,
        epsilon: float = 1e-4,
    ):
        self.server_lr = server_lr
        self.beta_1 = beta_1
        self.beta_2 = beta_2
        self.epsilon = epsilon
        
        self.m_t: Optional[List[np.ndarray]] = None
        self.v_t: Optional[List[np.ndarray]] = None
        self.t = 0
    
    def aggregate(
        self,
        parameters: List[List[np.ndarray]],
        weights: List[float],
        client_metrics: List[Dict],
    ) -> List[np.ndarray]:
        """Aggregate using FedOpt (server-side Adam)."""
        if not parameters:
            raise ValueError("No parameters to aggregate")
        
        # FedAvg as base
        base_aggregate = FedAvgAggregator().aggregate(parameters, weights, client_metrics)
        
        # Compute delta from base to weighted average
        if self.m_t is None:
            self.m_t = [np.zeros_like(p) for p in base_aggregate]
            self.v_t = [np.zeros_like(p) for p in base_aggregate]
        
        delta = [
            ba - p0
            for ba, p0 in zip(base_aggregate, parameters[0])
        ]
        
        self.t += 1
        
        # Update momentum and second moment
        self.m_t = [
            self.beta_1 * m + (1 - self.beta_1) * d
            for m, d in zip(self.m_t, delta)
        ]
        
        self.v_t = [
            self.beta_2 * v + (1 - self.beta_2) * (d ** 2)
            for v, d in zip(self.v_t, delta)
        ]
        
        # Bias correction
        m_hat = [m / (1 - self.beta_1 ** self.t) for m in self.m_t]
        v_hat = [v / (1 - self.beta_2 ** self.t) for v in self.v_t]
        
        # Apply update
        aggregated = [
            p0 + self.server_lr * m / (np.sqrt(v) + self.epsilon)
            for p0, m, v in zip(parameters[0], m_hat, v_hat)
        ]
        
        return aggregated


def create_aggregator(
    method: str = "fedavg",
    **kwargs,
) -> Aggregator:
    """Factory function to create aggregator.
    
    Args:
        method: Aggregation method ("fedavg", "fedprox", "scaffold", "fedopt")
        **kwargs: Additional arguments for the aggregator
    
    Returns:
        Aggregator instance
    """
    if method == "fedavg":
        return FedAvgAggregator()
    elif method == "fedprox":
        return FedProxAggregator(mu=kwargs.get("mu", 0.01))
    elif method == "scaffold":
        return SCAFFOLDAggregator()
    elif method == "fedopt":
        return FedOptAggregator(
            server_lr=kwargs.get("server_lr", 1.0),
            beta_1=kwargs.get("beta_1", 0.9),
            beta_2=kwargs.get("beta_2", 0.99),
        )
    else:
        raise ValueError(f"Unknown aggregation method: {method}")


if __name__ == "__main__":
    print("Available aggregators: fedavg, fedprox, scaffold, fedopt")
    
    # Example usage
    agg = create_aggregator("fedavg")
    print(f"Created aggregator: {type(agg).__name__}")