MEDfl 2.0.10.dev2__tar.gz → 2.0.10.dev3__tar.gz

{medfl-2.0.10.dev2 → medfl-2.0.10.dev3}/MEDfl/LearningManager/federated_dataset.py

@@ -11,6 +11,7 @@ class FederatedDataset:
         trainloaders: list,
         valloaders: list,
         testloaders: list,
+        feature_names: list[str] | None = None,
     ):
         """
         Represents a Federated Dataset.
@@ -30,6 +31,25 @@ class FederatedDataset:
         self.testloaders = testloaders
         self.size = len(self.trainloaders[0].dataset[0][0])
 
+        if feature_names is None:
+            feature_names = [
+                f"feature_{index}"
+                for index in range(self.size)
+            ]
+
+        if len(feature_names) != self.size:
+            raise ValueError(
+                "feature_names length must match the "
+                f"federated dataset size. Expected {self.size}, "
+                f"received {len(feature_names)}."
+            )
+
+        self.feature_names = list(
+            feature_names
+        )
+
+
+
         db_manager = DatabaseManager()
         db_manager.connect()
         self.eng = db_manager.get_connection()

{medfl-2.0.10.dev2 → medfl-2.0.10.dev3}/MEDfl/LearningManager/server.py

@@ -15,6 +15,12 @@ from .federated_dataset import FederatedDataset
 from .model import Model
 from .strategy import Strategy
 
+from .shap import (
+    SHAPConfig,
+    LocalPyTorchSHAPExplainer,
+    aggregate_federated_shap,
+)
+
 
 # ---------------------------------------------------------------------------
 # Logging setup — call this ONCE before starting the simulation.
@@ -64,7 +70,10 @@ class FlowerServer:
         num_clients: int,
         fed_dataset: FederatedDataset,
         diff_privacy: bool = False,
-        client_resources: Optional[Dict[str, float]] = {'num_cpus': 1, 'num_gpus': 0.0}
+        client_resources: Optional[Dict[str, float]] = {'num_cpus': 1, 'num_gpus': 0.0} , 
+        shap_config: Optional[
+        SHAPConfig
+    ] = None,
     ) -> None:
 
         self.device = torch.device("cpu")
@@ -77,6 +86,19 @@ class FlowerServer:
         self.strategy = strategy
         self.client_resources = client_resources
 
+        self.shap_config = (
+            shap_config
+            if shap_config is not None
+            else SHAPConfig(enabled=False)
+        )
+
+        self.shap_config.validate(
+            input_size=self.fed_dataset.size
+        )
+
+        self.federated_shap_result = None
+        self.local_shap_results = []
+
         setattr(self.strategy.strategy_object, "min_available_clients", self.num_clients)
         setattr(self.strategy.strategy_object, "initial_parameters", fl.common.ndarrays_to_parameters(self.params))
         setattr(self.strategy.strategy_object, "evaluate_fn", self.evaluate)
@@ -117,6 +139,319 @@ class FlowerServer:
         if hasattr(self.global_model.criterion, "pos_weight") and self.global_model.criterion.pos_weight is not None:
             self.global_model.criterion.pos_weight = self.global_model.criterion.pos_weight.detach().cpu()
         self.validate()
+    
+    def synchronize_final_global_model(
+        self,
+    ) -> None:
+        """
+        Copy the last Flower-aggregated parameters into
+        the MEDfl global model wrapper.
+        """
+
+        strategy_object = (
+            self.strategy.strategy_object
+        )
+
+        if strategy_object is None:
+            raise RuntimeError(
+                "The Flower strategy has not been created"
+            )
+
+        final_parameters = getattr(
+            strategy_object,
+            "latest_aggregated_parameters",
+            None,
+        )
+
+        if final_parameters is None:
+            # Your evaluate() method already records the latest
+            # centralized-evaluation parameter arrays.
+            final_arrays = getattr(
+                self,
+                "last_parameters",
+                None,
+            )
+
+            if final_arrays is None:
+                raise RuntimeError(
+                    "The final global model parameters "
+                    "could not be found"
+                )
+        else:
+            final_arrays = (
+                fl.common.parameters_to_ndarrays(
+                    final_parameters
+                )
+            )
+
+        self.global_model.set_parameters(
+            final_arrays
+        )
+
+        self.global_model.model.to(
+            self.device
+        )
+        self.global_model.model.eval()
+
+        logging.getLogger(
+            "MEDfl.server"
+        ).info(
+            "Final global parameters synchronized "
+            "for post-training operations"
+        )
+    
+    def _get_shap_loaders(
+        self,
+    ) -> list[DataLoader]:
+        """
+        Return one loader per client according to SHAPConfig.
+        """
+
+        split = self.shap_config.data_split
+
+        if split == "train":
+            return self.fed_dataset.trainloaders
+
+        if split == "validation":
+            return self.fed_dataset.valloaders
+
+        if split == "test":
+            return self.fed_dataset.testloaders
+
+        raise ValueError(
+            f"Unsupported SHAP data split: {split}"
+        )
+    
+    def calculate_local_shap_results(
+        self,
+    ) -> list[dict]:
+        """
+        Calculate local SHAP summaries for every client partition.
+
+        This method runs only after federated training has finished.
+        """
+
+        if not self.shap_config.enabled:
+            return []
+
+        shap_logger = logging.getLogger(
+            "MEDfl.shap"
+        )
+
+        loaders = self._get_shap_loaders()
+
+        feature_names = (
+            self.shap_config.feature_names
+            if self.shap_config.feature_names
+            is not None
+            else self.fed_dataset.feature_names
+        )
+
+        if len(feature_names) != self.fed_dataset.size:
+            raise ValueError(
+                "SHAP feature names do not match the "
+                "federated dataset input size"
+            )
+
+        local_explainer = (
+            LocalPyTorchSHAPExplainer(
+                explainer_type=(
+                    self.shap_config.explainer
+                ),
+                background_size=(
+                    self.shap_config.background_size
+                ),
+                explanation_size=(
+                    self.shap_config.explanation_size
+                ),
+                random_seed=(
+                    self.shap_config.random_seed
+                ),
+                clipping_value=(
+                    self.shap_config.clipping_value
+                ),
+            )
+        )
+
+        local_results: list[dict] = []
+
+        for client_index, loader in enumerate(
+            loaders
+        ):
+            client_id = str(client_index)
+
+            if loader is None:
+                shap_logger.warning(
+                    "Skipping SHAP client %s: loader is None",
+                    client_id,
+                )
+                continue
+
+            available_samples = len(
+                loader.dataset
+            )
+
+            if available_samples == 0:
+                shap_logger.warning(
+                    "Skipping SHAP client %s: empty dataset",
+                    client_id,
+                )
+                continue
+
+            shap_logger.info(
+                "Starting local SHAP | client=%s | "
+                "available_samples=%d",
+                client_id,
+                available_samples,
+            )
+
+            try:
+                # The model is used for inference only.
+                # No optimizer or training method is called.
+                local_result = (
+                    local_explainer.calculate(
+                        model=self.global_model.model,
+                        loader=loader,
+                        device=self.device,
+                    )
+                )
+
+                explained_samples = int(
+                    local_result["sample_count"]
+                )
+
+                if (
+                    explained_samples
+                    < self.shap_config.minimum_samples
+                ):
+                    shap_logger.warning(
+                        "Skipping client %s because only %d "
+                        "samples were explained; minimum=%d",
+                        client_id,
+                        explained_samples,
+                        self.shap_config.minimum_samples,
+                    )
+                    continue
+
+                local_result[
+                    "client_id"
+                ] = client_id
+
+                local_result[
+                    "feature_names"
+                ] = list(feature_names)
+
+                local_result[
+                    "available_samples"
+                ] = available_samples
+
+                local_results.append(
+                    local_result
+                )
+
+                shap_logger.info(
+                    "Local SHAP completed | client=%s | "
+                    "explained_samples=%d",
+                    client_id,
+                    explained_samples,
+                )
+
+            except Exception as exc:
+                shap_logger.exception(
+                    "Local SHAP failed for client %s: %s",
+                    client_id,
+                    exc,
+                )
+
+        if not local_results:
+            raise RuntimeError(
+                "SHAP was enabled, but no client produced "
+                "a valid explanation result"
+            )
+
+        self.local_shap_results = (
+            local_results
+        )
+
+        return local_results
+    
+    def calculate_federated_shap(
+        self,
+    ) -> dict:
+        """
+        Run the complete post-training SHAP phase.
+        """
+
+        shap_logger = logging.getLogger(
+            "MEDfl.shap"
+        )
+
+        if not self.shap_config.enabled:
+            return {
+                "status": "disabled",
+            }
+
+        shap_logger.info(
+            "========================================"
+        )
+        shap_logger.info(
+            "Starting post-training federated SHAP"
+        )
+        shap_logger.info(
+            "Explainer: %s",
+            self.shap_config.explainer,
+        )
+        shap_logger.info(
+            "Data split: %s",
+            self.shap_config.data_split,
+        )
+        shap_logger.info(
+            "========================================"
+        )
+
+        local_results = (
+            self.calculate_local_shap_results()
+        )
+
+        federated_result = (
+            aggregate_federated_shap(
+                client_results=local_results,
+                include_client_results=(
+                    self.shap_config
+                    .include_client_results
+                ),
+            )
+        )
+
+        federated_result["explainer"] = (
+            self.shap_config.explainer
+        )
+        federated_result["data_split"] = (
+            self.shap_config.data_split
+        )
+        federated_result["background_size"] = (
+            self.shap_config.background_size
+        )
+        federated_result["requested_explanation_size"] = (
+            self.shap_config.explanation_size
+        )
+
+        self.federated_shap_result = (
+            federated_result
+        )
+
+        shap_logger.info(
+            "Federated SHAP completed | clients=%d | "
+            "samples=%d",
+            federated_result[
+                "participating_clients"
+            ],
+            federated_result[
+                "explained_samples"
+            ],
+        )
+
+        return federated_result
 
     def validate(self) -> None:
         if not isinstance(self.global_model, Model):
@@ -249,4 +584,32 @@ class FlowerServer:
             ray_init_args=ray_init_args,
             client_resources=self.client_resources
         )
+    # Training has now completely finished.
+        self.synchronize_final_global_model()
+
+        if self.shap_config.enabled:
+            try:
+                self.federated_shap_result = (
+                    self.calculate_federated_shap()
+                )
+
+            except Exception as exc:
+                logging.getLogger(
+                    "MEDfl.shap"
+                ).exception(
+                    "Post-training federated SHAP failed: %s",
+                    exc,
+                )
+
+                # Important: a SHAP failure does not mean
+                # federated training failed.
+                self.federated_shap_result = {
+                    "status": "failed",
+                    "error": str(exc),
+                }
+        else:
+            self.federated_shap_result = {
+                "status": "disabled",
+            }
+
         return history

medfl-2.0.10.dev3/MEDfl/LearningManager/shap/__init__.py

@@ -0,0 +1,9 @@
+from .config import SHAPConfig
+from .explainer import LocalPyTorchSHAPExplainer
+from .aggregation import aggregate_federated_shap
+
+__all__ = [
+    "SHAPConfig",
+    "LocalPyTorchSHAPExplainer",
+    "aggregate_federated_shap",
+]

medfl-2.0.10.dev3/MEDfl/LearningManager/shap/aggregation.py

@@ -0,0 +1,249 @@
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+
+def _validate_client_results(
+    client_results: list[dict],
+) -> tuple[list[str], int]:
+    if not client_results:
+        raise ValueError(
+            "No client SHAP results were provided"
+        )
+
+    reference_features = client_results[0][
+        "feature_names"
+    ]
+
+    feature_count = len(reference_features)
+
+    required_array_keys = {
+        "abs_sum",
+        "signed_sum",
+        "square_sum",
+        "positive_count",
+        "negative_count",
+        "zero_count",
+    }
+
+    for result in client_results:
+        if result["feature_names"] != reference_features:
+            raise ValueError(
+                "Clients used different feature names or orders"
+            )
+
+        if int(result["sample_count"]) <= 0:
+            raise ValueError(
+                f"Client {result.get('client_id')} returned an "
+                "invalid SHAP sample count"
+            )
+
+        for key in required_array_keys:
+            if key not in result:
+                raise ValueError(
+                    f"Client result is missing '{key}'"
+                )
+
+            values = np.asarray(result[key])
+
+            if values.shape != (feature_count,):
+                raise ValueError(
+                    f"Client {result.get('client_id')} returned "
+                    f"invalid shape {values.shape} for {key}. "
+                    f"Expected ({feature_count},)."
+                )
+
+    return reference_features, feature_count
+
+
+def aggregate_federated_shap(
+    client_results: list[dict],
+    include_client_results: bool = True,
+) -> dict[str, Any]:
+    """
+    Aggregate additive local SHAP statistics.
+
+    The global mean absolute SHAP is equivalent to calculating:
+
+        sum of all |SHAP values| / number of all explained samples
+    """
+
+    feature_names, feature_count = (
+        _validate_client_results(client_results)
+    )
+
+    total_samples = sum(
+        int(result["sample_count"])
+        for result in client_results
+    )
+
+    abs_sum = np.zeros(
+        feature_count,
+        dtype=np.float64,
+    )
+    signed_sum = np.zeros(
+        feature_count,
+        dtype=np.float64,
+    )
+    square_sum = np.zeros(
+        feature_count,
+        dtype=np.float64,
+    )
+    positive_count = np.zeros(
+        feature_count,
+        dtype=np.float64,
+    )
+    negative_count = np.zeros(
+        feature_count,
+        dtype=np.float64,
+    )
+    zero_count = np.zeros(
+        feature_count,
+        dtype=np.float64,
+    )
+
+    for result in client_results:
+        abs_sum += np.asarray(
+            result["abs_sum"],
+            dtype=np.float64,
+        )
+        signed_sum += np.asarray(
+            result["signed_sum"],
+            dtype=np.float64,
+        )
+        square_sum += np.asarray(
+            result["square_sum"],
+            dtype=np.float64,
+        )
+        positive_count += np.asarray(
+            result["positive_count"],
+            dtype=np.float64,
+        )
+        negative_count += np.asarray(
+            result["negative_count"],
+            dtype=np.float64,
+        )
+        zero_count += np.asarray(
+            result["zero_count"],
+            dtype=np.float64,
+        )
+
+    mean_abs = abs_sum / total_samples
+    mean_signed = signed_sum / total_samples
+    mean_square = square_sum / total_samples
+
+    # Var(X) = E(X²) - E(X)²
+    variance = np.maximum(
+        mean_square - np.square(mean_signed),
+        0.0,
+    )
+
+    standard_deviation = np.sqrt(variance)
+
+    features = []
+
+    for index, feature_name in enumerate(
+        feature_names
+    ):
+        features.append(
+            {
+                "feature": feature_name,
+                "feature_index": index,
+                "mean_abs_shap": float(
+                    mean_abs[index]
+                ),
+                "mean_signed_shap": float(
+                    mean_signed[index]
+                ),
+                "shap_std": float(
+                    standard_deviation[index]
+                ),
+                "positive_rate": float(
+                    positive_count[index]
+                    / total_samples
+                ),
+                "negative_rate": float(
+                    negative_count[index]
+                    / total_samples
+                ),
+                "zero_rate": float(
+                    zero_count[index]
+                    / total_samples
+                ),
+            }
+        )
+
+    # Highest global importance first.
+    features.sort(
+        key=lambda item: item["mean_abs_shap"],
+        reverse=True,
+    )
+
+    final_result: dict[str, Any] = {
+        "status": "completed",
+        "participating_clients": len(
+            client_results
+        ),
+        "explained_samples": int(
+            total_samples
+        ),
+        "feature_importance": features,
+    }
+
+    if include_client_results:
+        client_summaries = []
+
+        for result in client_results:
+            sample_count = int(
+                result["sample_count"]
+            )
+
+            client_mean_abs = (
+                np.asarray(
+                    result["abs_sum"],
+                    dtype=np.float64,
+                )
+                / sample_count
+            )
+
+            local_features = [
+                {
+                    "feature": feature_name,
+                    "feature_index": index,
+                    "mean_abs_shap": float(
+                        client_mean_abs[index]
+                    ),
+                }
+                for index, feature_name in enumerate(
+                    feature_names
+                )
+            ]
+
+            local_features.sort(
+                key=lambda item: item[
+                    "mean_abs_shap"
+                ],
+                reverse=True,
+            )
+
+            client_summaries.append(
+                {
+                    "client_id": result[
+                        "client_id"
+                    ],
+                    "explained_samples": (
+                        sample_count
+                    ),
+                    "feature_importance": (
+                        local_features
+                    ),
+                }
+            )
+
+        final_result["client_results"] = (
+            client_summaries
+        )
+
+    return final_result