MEDfl 2.0.10.dev0__tar.gz → 2.0.10.dev2__tar.gz

medfl-2.0.10.dev2/MEDfl/LearningManager/xgboost/__init__.py

@@ -0,0 +1,15 @@
+"""
+Federated XGBoost backend for MEDfl simulation mode.
+"""
+
+from .client import XGBoostSimulationClient
+from .config import XGBoostConfig
+from .server import XGBoostSimulationServer
+from .strategy import XGBoostSimulationStrategy
+
+__all__ = [
+    "XGBoostSimulationClient",
+    "XGBoostConfig",
+    "XGBoostSimulationServer",
+    "XGBoostSimulationStrategy",
+]

medfl-2.0.10.dev2/MEDfl/LearningManager/xgboost/client.py

@@ -0,0 +1,328 @@
+#!/usr/bin/env python3
+
+import logging
+import sys
+from typing import Any, Dict, List, Optional
+
+import flwr as fl
+import xgboost as xgb
+from flwr.common import (
+    Code,
+    EvaluateIns,
+    EvaluateRes,
+    FitIns,
+    FitRes,
+    GetParametersIns,
+    GetParametersRes,
+    GetPropertiesIns,
+    GetPropertiesRes,
+    Status,
+)
+from torch.utils.data import DataLoader
+
+from .data import SimulationXGBoostDataModule
+from .metrics import evaluate_predictions
+from .serialization import (
+    booster_to_parameters,
+    parameters_to_booster,
+)
+
+
+def _get_logger() -> logging.Logger:
+    logger = logging.getLogger(
+        "MEDfl.simulation.xgboost.client"
+    )
+
+    if not logger.handlers:
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setFormatter(
+            logging.Formatter(
+                "[%(asctime)s][%(name)s][%(levelname)s] "
+                "%(message)s",
+                datefmt="%H:%M:%S",
+            )
+        )
+
+        logger.addHandler(handler)
+        logger.setLevel(logging.INFO)
+        logger.propagate = False
+
+    return logger
+
+
+class XGBoostSimulationClient(fl.client.Client):
+    def __init__(
+        self,
+        cid: str,
+        trainloader: DataLoader,
+        valloader: Optional[DataLoader] = None,
+        testloader: Optional[DataLoader] = None,
+        task: str = "binary",
+        xgb_params: Optional[Dict[str, Any]] = None,
+        local_num_boost_round: int = 10,
+        threshold: float = 0.5,
+        feature_names: Optional[List[str]] = None,
+    ):
+        self.cid = str(cid)
+        self.task = task
+        self.xgb_params = xgb_params or {}
+        self.local_num_boost_round = int(
+            local_num_boost_round
+        )
+        self.threshold = float(threshold)
+        self.feature_names = feature_names
+
+        self.log = _get_logger()
+        self.booster: Optional[xgb.Booster] = None
+
+        data_module = SimulationXGBoostDataModule(
+            trainloader=trainloader,
+            valloader=valloader,
+            testloader=testloader,
+            feature_names=feature_names,
+        )
+
+        self.dtrain = data_module.dtrain
+        self.dval = data_module.dval
+        self.dtest = data_module.dtest
+        self.metadata = data_module.metadata
+
+        self._validate()
+
+        self.log.info(
+            "[Client %s] initialized: "
+            "train=%s, val=%s, test=%s, features=%s",
+            self.cid,
+            self.metadata["num_train"],
+            self.metadata["num_val"],
+            self.metadata["num_test"],
+            self.metadata["num_features"],
+        )
+
+    def _validate(self) -> None:
+        if self.task not in {
+            "binary",
+            "regression",
+            "multiclass",
+        }:
+            raise ValueError(
+                f"Unsupported task: {self.task}"
+            )
+
+        if self.dtrain is None:
+            raise ValueError(
+                f"Client {self.cid} has no training data."
+            )
+
+        if self.local_num_boost_round <= 0:
+            raise ValueError(
+                "local_num_boost_round must be greater than zero."
+            )
+
+    def get_parameters(
+        self,
+        ins: GetParametersIns,
+    ) -> GetParametersRes:
+        return GetParametersRes(
+            status=Status(
+                code=Code.OK,
+                message="Success",
+            ),
+            parameters=booster_to_parameters(
+                self.booster
+            ),
+        )
+
+    def fit(self, ins: FitIns) -> FitRes:
+        self.log.info(
+            "[Client %s] starting local XGBoost fit",
+            self.cid,
+        )
+
+        received_booster = parameters_to_booster(
+            ins.parameters
+        )
+
+        if received_booster is not None:
+            self.booster = received_booster
+
+        evals = []
+
+        if self.dval is not None:
+            evals.append(
+                (self.dval, "validation")
+            )
+
+        self.booster = xgb.train(
+            params=self.xgb_params,
+            dtrain=self.dtrain,
+            num_boost_round=self.local_num_boost_round,
+            xgb_model=self.booster,
+            evals=evals,
+            verbose_eval=False,
+        )
+
+        train_predictions = self.booster.predict(
+            self.dtrain
+        )
+
+        train_metrics = evaluate_predictions(
+            task=self.task,
+            y_true=self.dtrain.get_label(),
+            y_pred_or_prob=train_predictions,
+            threshold=self.threshold,
+            num_features=int(
+                self.metadata["num_features"]
+            ),
+        )
+
+        metrics = {
+            "backend": "xgboost",
+            "task": self.task,
+            "cid": self.cid,
+            "num_train": int(
+                self.metadata["num_train"]
+            ),
+        }
+
+        for key, value in train_metrics.items():
+            metrics[f"train_{key}"] = value
+
+        self.log.info(
+            "[Client %s] fit completed: %s",
+            self.cid,
+            metrics,
+        )
+
+        return FitRes(
+            status=Status(
+                code=Code.OK,
+                message="Success",
+            ),
+            parameters=booster_to_parameters(
+                self.booster
+            ),
+            num_examples=int(
+                self.metadata["num_train"]
+            ),
+            metrics=metrics,
+        )
+
+    def evaluate(
+        self,
+        ins: EvaluateIns,
+    ) -> EvaluateRes:
+        received_booster = parameters_to_booster(
+            ins.parameters
+        )
+
+        if received_booster is not None:
+            self.booster = received_booster
+
+        evaluation_matrix = (
+            self.dtest
+            if self.dtest is not None
+            else self.dval
+        )
+
+        if (
+            self.booster is None
+            or evaluation_matrix is None
+        ):
+            return EvaluateRes(
+                status=Status(
+                    code=Code.EVALUATE_NOT_IMPLEMENTED,
+                    message=(
+                        "No booster or evaluation data."
+                    ),
+                ),
+                loss=0.0,
+                num_examples=0,
+                metrics={},
+            )
+
+        predictions = self.booster.predict(
+            evaluation_matrix
+        )
+
+        eval_metrics = evaluate_predictions(
+            task=self.task,
+            y_true=evaluation_matrix.get_label(),
+            y_pred_or_prob=predictions,
+            threshold=self.threshold,
+            num_features=int(
+                self.metadata["num_features"]
+            ),
+        )
+
+        metrics = {
+            "backend": "xgboost",
+            "task": self.task,
+            "cid": self.cid,
+        }
+
+        for key, value in eval_metrics.items():
+            metrics[f"eval_{key}"] = value
+
+        if self.task == "binary":
+            loss = float(
+                eval_metrics.get("logloss", 0.0)
+            )
+        elif self.task == "regression":
+            loss = float(
+                eval_metrics.get("rmse", 0.0)
+            )
+        else:
+            loss = float(
+                eval_metrics.get("mlogloss", 0.0)
+            )
+
+        self.log.info(
+            "[Client %s] evaluation completed: "
+            "loss=%s metrics=%s",
+            self.cid,
+            loss,
+            metrics,
+        )
+
+        return EvaluateRes(
+            status=Status(
+                code=Code.OK,
+                message="Success",
+            ),
+            loss=loss,
+            num_examples=int(
+                evaluation_matrix.num_row()
+            ),
+            metrics=metrics,
+        )
+
+    def get_properties(
+        self,
+        ins: GetPropertiesIns,
+    ) -> GetPropertiesRes:
+        properties = {
+            "backend": "xgboost",
+            "cid": self.cid,
+            "task": self.task,
+            "num_train": int(
+                self.metadata["num_train"]
+            ),
+            "num_val": int(
+                self.metadata["num_val"]
+            ),
+            "num_test": int(
+                self.metadata["num_test"]
+            ),
+            "num_features": int(
+                self.metadata["num_features"]
+            ),
+        }
+
+        return GetPropertiesRes(
+            status=Status(
+                code=Code.OK,
+                message="Success",
+            ),
+            properties=properties,
+        )

medfl-2.0.10.dev2/MEDfl/LearningManager/xgboost/config.py

@@ -0,0 +1,63 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class XGBoostConfig:
+    task: str = "binary"
+    mode: str = "bagging"
+    local_num_boost_round: int = 10
+    params: Dict[str, Any] = field(default_factory=dict)
+    num_classes: Optional[int] = None
+
+    def build_params(self) -> Dict[str, Any]:
+        if self.task == "binary":
+            defaults = {
+                "objective": "binary:logistic",
+                "eval_metric": "auc",
+                "tree_method": "hist",
+                "max_depth": 6,
+                "eta": 0.1,
+                "subsample": 1.0,
+                "colsample_bytree": 1.0,
+            }
+
+        elif self.task == "regression":
+            defaults = {
+                "objective": "reg:squarederror",
+                "eval_metric": "rmse",
+                "tree_method": "hist",
+                "max_depth": 6,
+                "eta": 0.1,
+                "subsample": 1.0,
+                "colsample_bytree": 1.0,
+            }
+
+        elif self.task == "multiclass":
+            if self.num_classes is None or self.num_classes < 2:
+                raise ValueError(
+                    "num_classes must be provided and >= 2 "
+                    "for multiclass XGBoost."
+                )
+
+            defaults = {
+                "objective": "multi:softprob",
+                "num_class": self.num_classes,
+                "eval_metric": "mlogloss",
+                "tree_method": "hist",
+                "max_depth": 6,
+                "eta": 0.1,
+                "subsample": 1.0,
+                "colsample_bytree": 1.0,
+            }
+
+        else:
+            raise ValueError(
+                f"Unsupported XGBoost task: {self.task}. "
+                "Expected binary, regression, or multiclass."
+            )
+
+        merged = defaults.copy()
+        merged.update(self.params or {})
+
+        return merged

medfl-2.0.10.dev2/MEDfl/LearningManager/xgboost/data.py

@@ -0,0 +1,122 @@
+from typing import List, Optional, Tuple
+
+import numpy as np
+import torch
+import xgboost as xgb
+from torch.utils.data import DataLoader
+
+
+def dataloader_to_numpy(
+    loader: DataLoader,
+) -> Tuple[np.ndarray, np.ndarray]:
+    if loader is None:
+        raise ValueError("DataLoader cannot be None.")
+
+    x_batches = []
+    y_batches = []
+
+    for x_batch, y_batch in loader:
+        if isinstance(x_batch, torch.Tensor):
+            x_batch = x_batch.detach().cpu().numpy()
+
+        if isinstance(y_batch, torch.Tensor):
+            y_batch = y_batch.detach().cpu().numpy()
+
+        x_batch = np.asarray(x_batch)
+        y_batch = np.asarray(y_batch).reshape(-1)
+
+        x_batches.append(x_batch)
+        y_batches.append(y_batch)
+
+    if not x_batches:
+        raise ValueError(
+            "Cannot convert an empty DataLoader to DMatrix."
+        )
+
+    x = np.concatenate(x_batches, axis=0)
+    y = np.concatenate(y_batches, axis=0)
+
+    if x.ndim != 2:
+        raise ValueError(
+            "XGBoost tabular simulation expects a two-dimensional "
+            f"feature matrix. Received shape={x.shape}."
+        )
+
+    return x, y
+
+
+def dataloader_to_dmatrix(
+    loader: Optional[DataLoader],
+    feature_names: Optional[List[str]] = None,
+) -> Optional[xgb.DMatrix]:
+    if loader is None:
+        return None
+
+    try:
+        if len(loader.dataset) == 0:
+            return None
+    except TypeError:
+        pass
+
+    x, y = dataloader_to_numpy(loader)
+
+    if feature_names is not None:
+        if len(feature_names) != x.shape[1]:
+            raise ValueError(
+                "feature_names length does not match the number "
+                f"of features: {len(feature_names)} != {x.shape[1]}"
+            )
+
+    return xgb.DMatrix(
+        data=x,
+        label=y,
+        feature_names=feature_names,
+    )
+
+
+class SimulationXGBoostDataModule:
+    def __init__(
+        self,
+        trainloader: DataLoader,
+        valloader: Optional[DataLoader] = None,
+        testloader: Optional[DataLoader] = None,
+        feature_names: Optional[List[str]] = None,
+    ):
+        self.dtrain = dataloader_to_dmatrix(
+            trainloader,
+            feature_names,
+        )
+
+        self.dval = dataloader_to_dmatrix(
+            valloader,
+            feature_names,
+        )
+
+        self.dtest = dataloader_to_dmatrix(
+            testloader,
+            feature_names,
+        )
+
+        self.metadata = {
+            "num_train": (
+                self.dtrain.num_row()
+                if self.dtrain is not None
+                else 0
+            ),
+            "num_val": (
+                self.dval.num_row()
+                if self.dval is not None
+                else 0
+            ),
+            "num_test": (
+                self.dtest.num_row()
+                if self.dtest is not None
+                else 0
+            ),
+            "num_features": (
+                self.dtrain.num_col()
+                if self.dtrain is not None
+                else 0
+            ),
+            "feature_names": feature_names or [],
+        }

medfl-2.0.10.dev2/MEDfl/LearningManager/xgboost/metrics.py

@@ -0,0 +1,192 @@
+import math
+from typing import Any, Dict
+
+import numpy as np
+from sklearn.metrics import (
+    accuracy_score,
+    f1_score,
+    log_loss,
+    mean_absolute_error,
+    mean_squared_error,
+    r2_score,
+    roc_auc_score,
+)
+
+
+def _safe_float(value: Any) -> float:
+    try:
+        value = float(value)
+
+        if math.isnan(value) or math.isinf(value):
+            return float("nan")
+
+        return value
+    except Exception:
+        return float("nan")
+
+
+def binary_metrics(
+    y_true,
+    y_prob,
+    threshold: float = 0.5,
+) -> Dict[str, float]:
+    y_true = np.asarray(y_true)
+    y_prob = np.asarray(y_prob)
+    y_pred = (y_prob >= threshold).astype(int)
+
+    metrics = {
+        "accuracy": _safe_float(
+            accuracy_score(y_true, y_pred)
+        ),
+    }
+
+    try:
+        metrics["auc"] = _safe_float(
+            roc_auc_score(y_true, y_prob)
+        )
+    except Exception:
+        metrics["auc"] = float("nan")
+
+    try:
+        metrics["logloss"] = _safe_float(
+            log_loss(y_true, y_prob)
+        )
+    except Exception:
+        metrics["logloss"] = float("nan")
+
+    return metrics
+
+
+def regression_metrics(
+    y_true,
+    y_pred,
+    num_features: int = 0,
+) -> Dict[str, float]:
+    y_true = np.asarray(y_true)
+    y_pred = np.asarray(y_pred)
+
+    rmse = mean_squared_error(
+        y_true,
+        y_pred,
+        squared=False,
+    )
+    mae = mean_absolute_error(y_true, y_pred)
+    r2 = r2_score(y_true, y_pred)
+
+    n = len(y_true)
+    p = int(num_features)
+
+    if n > p + 1:
+        adjusted_r2 = (
+            1.0
+            - (1.0 - r2)
+            * ((n - 1.0) / (n - p - 1.0))
+        )
+    else:
+        adjusted_r2 = float("nan")
+
+    return {
+        "rmse": _safe_float(rmse),
+        "mae": _safe_float(mae),
+        "r2": _safe_float(r2),
+        "adjusted_r2": _safe_float(adjusted_r2),
+    }
+
+
+def multiclass_metrics(
+    y_true,
+    y_prob,
+) -> Dict[str, float]:
+    y_true = np.asarray(y_true)
+    y_prob = np.asarray(y_prob)
+    y_pred = np.argmax(y_prob, axis=1)
+
+    metrics = {
+        "accuracy": _safe_float(
+            accuracy_score(y_true, y_pred)
+        ),
+    }
+
+    try:
+        metrics["macro_f1"] = _safe_float(
+            f1_score(
+                y_true,
+                y_pred,
+                average="macro",
+            )
+        )
+    except Exception:
+        metrics["macro_f1"] = float("nan")
+
+    try:
+        metrics["mlogloss"] = _safe_float(
+            log_loss(y_true, y_prob)
+        )
+    except Exception:
+        metrics["mlogloss"] = float("nan")
+
+    return metrics
+
+
+def evaluate_predictions(
+    task: str,
+    y_true,
+    y_pred_or_prob,
+    threshold: float = 0.5,
+    num_features: int = 0,
+) -> Dict[str, float]:
+    if task == "binary":
+        return binary_metrics(
+            y_true,
+            y_pred_or_prob,
+            threshold=threshold,
+        )
+
+    if task == "regression":
+        return regression_metrics(
+            y_true,
+            y_pred_or_prob,
+            num_features=num_features,
+        )
+
+    if task == "multiclass":
+        return multiclass_metrics(
+            y_true,
+            y_pred_or_prob,
+        )
+
+    raise ValueError(f"Unsupported task: {task}")
+
+
+def aggregate_weighted_metrics(results):
+    aggregated = {}
+
+    keys = set()
+    for _, metrics in results:
+        keys.update(metrics.keys())
+
+    for key in keys:
+        weighted_sum = 0.0
+        contributing_examples = 0
+
+        for num_examples, metrics in results:
+            if key not in metrics:
+                continue
+
+            try:
+                value = float(metrics[key])
+            except Exception:
+                continue
+
+            if math.isnan(value) or math.isinf(value):
+                continue
+
+            weighted_sum += num_examples * value
+            contributing_examples += num_examples
+
+        if contributing_examples > 0:
+            aggregated[key] = (
+                weighted_sum / contributing_examples
+            )
+
+    return aggregated