API reference

High-level estimator

A compact, model-independent entry point.

Parameters use user-facing names while mapping onto the backward-compatible :class:FSAConfig fields. Advanced users may pass an existing config instead.

Source code in federated_survival/api.py

class FederatedSurvival:
    """A compact, model-independent entry point.

    Parameters use user-facing names while mapping onto the backward-compatible
    :class:`FSAConfig` fields.  Advanced users may pass an existing ``config``
    instead.
    """

    def __init__(
        self,
        model: str = "DeepSurv",
        *,
        protocol: str = "FedAvg",
        protocol_params: Optional[Dict[str, Any]] = None,
        config: Optional[FSAConfig] = None,
        n_clients: int = 5,
        global_rounds: int = 20,
        local_steps: int = 1,
        batch_size: int = 32,
        learning_rate: float = 1e-3,
        client_fraction: float = 1.0,
        split_type: str = "iid",
        split_alpha: float = 0.5,
        test_size: float = 0.2,
        random_state: int = 42,
        **config_overrides,
    ) -> None:
        if config is not None and (config_overrides or protocol_params):
            raise ValueError(
                "protocol_params/config_overrides cannot be combined with an existing config"
            )
        self.split_type = split_type
        self.split_alpha = split_alpha
        self.test_size = test_size
        if config is None:
            values = {
                "model_type": model,
                "federated_protocol": protocol,
                "num_clients": n_clients,
                "global_epochs": global_rounds,
                "local_epochs": local_steps,
                "batch_size": batch_size,
                "learning_rate": learning_rate,
                "client_sample_ratio": client_fraction,
                "split_method": split_type,
                "split_alpha": split_alpha,
                "test_size": test_size,
                "random_seed": random_state,
                "prediction_validation": "raise",
            }
            if protocol_params:
                parameter_aliases = {
                    "mu": "proximal_mu",
                    "optimizer": "server_optimizer",
                    "server_lr": "server_learning_rate",
                    "beta1": "server_beta1",
                    "beta2": "server_beta2",
                    "tau": "server_tau",
                    "max_iter": "webdisco_max_iter",
                    "tolerance": "webdisco_tolerance",
                    "l2": "webdisco_l2",
                    "ridge": "webdisco_ridge",
                    "max_step_norm": "webdisco_max_step_norm",
                }
                for key, value in protocol_params.items():
                    values[parameter_aliases.get(key, key)] = value
            values.update(config_overrides)
            if model == "CoxPH" and "num_nodes" not in values:
                values["num_nodes"] = ()
            config = FSAConfig(**values)
        else:
            self.split_type = config.split_method
            self.split_alpha = config.split_alpha
            self.test_size = config.test_size
        self.config = config
        self.runner_: Optional[FSARunner] = None
        self.history_: Optional[Dict[str, Any]] = None
        self.feature_names_: Optional[Sequence[str]] = None
        self.dataset_: Optional[DataSet] = None

    @property
    def available_models(self):
        return available_model_adapters()

    @property
    def available_protocols(self):
        return available_federated_protocols()

    def fit(
        self,
        data,
        durations=None,
        events=None,
        *,
        duration_col: str = "time",
        event_col: str = "status",
        client_col: Optional[str] = None,
        augmentation: Optional[str] = None,
    ) -> "FederatedSurvival":
        """Fit from a package ``DataSet``, DataFrame, or NumPy arrays.

        A DataFrame may optionally contain an institution column.  Otherwise
        the requested synthetic partition strategy is applied.  The feature
        count is always inferred from the supplied data.
        """

        dataset = self._prepare_dataset(
            data,
            durations,
            events,
            duration_col=duration_col,
            event_col=event_col,
            client_col=client_col,
        )
        first_x = next(iter(dataset.clients_set.values()))[0]
        self.config.n_features = int(first_x.shape[1])
        self.config.num_clients = len(dataset.clients_set)

        runner = FSARunner(self.config)
        if augmentation is None:
            history = runner.run(dataset, type="raw")
        else:
            normalized = augmentation.upper()
            if normalized not in {"MVAEC", "MVAES"}:
                raise ValueError("augmentation must be None, 'MVAEC', or 'MVAES'")
            history = runner.run(dataset, type="raw_aug", aug_method=normalized)

        self.runner_ = runner
        self.history_ = history
        self.dataset_ = dataset
        return self

    def predict_survival(self, features) -> pd.DataFrame:
        """Return a time-by-sample survival-probability DataFrame."""
        self._require_fitted()
        array = self._coerce_features(features)
        return self.runner_.predict_survival(array)

    def evaluate(self, features, durations, events) -> Dict[str, Any]:
        """Compute time-dependent concordance and IPCW integrated Brier score."""
        survival = self.predict_survival(features)
        duration_values = np.asarray(durations, dtype=float).reshape(-1)
        event_values = np.asarray(events, dtype=float).reshape(-1)
        if survival.shape[1] != len(duration_values) or len(events) != len(duration_values):
            raise ValueError("features, durations, and events must contain the same samples")
        evaluator = EvalSurv(
            survival,
            duration_values,
            event_values,
            censor_surv="km",
        )
        grid = evaluation_time_grid(
            duration_values,
            survival.index.values,
            self.config.evaluation_quantiles,
        )
        diagnostics = self.runner_.last_prediction_diagnostics_
        return {
            "c_index": float(evaluator.concordance_td()),
            "ibs": float(evaluator.integrated_brier_score(grid)),
            "prediction_diagnostics": asdict(diagnostics),
        }

    def get_run_summary(self) -> Dict[str, Any]:
        """Return serialisable configuration and final metrics."""
        self._require_fitted()
        final = {}
        for name in ("train_Cindex", "train_IBS", "test_Cindex", "test_IBS"):
            values = self.history_.get(name, [])
            final[name] = float(values[-1]) if values else None
        return {
            "model": self.config.model_type,
            "protocol": self.config.federated_protocol,
            "n_clients": self.config.num_clients,
            "n_features": self.config.n_features,
            "global_rounds": self.config.global_epochs,
            "local_steps": self.config.local_epochs,
            "batch_size": self.config.batch_size,
            "final_metrics": final,
        }

    def get_privacy_info(self) -> Dict[str, Any]:
        """Report configured perturbation settings, not an achieved DP budget."""
        self._require_fitted()
        config = self.config
        if not config.use_differential_privacy:
            return {"enabled": False, "formal_accounting_available": False}
        from .core.differential_privacy import describe_noise_knobs

        knobs = describe_noise_knobs(config.dp_mechanism)
        return {
            "enabled": True,
            "mechanism": config.dp_mechanism,
            "configured_epsilon": config.dp_epsilon,
            "configured_delta": config.dp_delta if config.dp_mechanism == "gaussian" else None,
            "configured_sensitivity": config.dp_sensitivity,
            "configured_noise_multiplier": (
                config.dp_noise_multiplier if config.dp_mechanism == "gaussian" else None
            ),
            "configured_clip_norm": (
                config.dp_clip_norm if config.dp_mechanism in ("gaussian", "laplace") else None
            ),
            "formal_accounting_available": False,
            "privacy_scope": "experimental clipped client-update perturbation",
            # The Gaussian scale follows dp_noise_multiplier and the Laplace
            # scale follows dp_epsilon, so the reported epsilon is not the
            # noise control under Gaussian.  Naming the driver and the ignored
            # fields keeps that from being misread.
            "noise_driver": knobs["noise_knob"],
            "inactive_parameters": list(knobs["inert_knobs"]),
        }

    def _prepare_dataset(
        self,
        data,
        durations,
        events,
        *,
        duration_col,
        event_col,
        client_col,
    ) -> DataSet:
        if isinstance(data, DataSet):
            first_x = next(iter(data.clients_set.values()))[0]
            self.feature_names_ = tuple(f"x{index + 1}" for index in range(first_x.shape[1]))
            return data

        if isinstance(data, pd.DataFrame):
            frame = data.copy()
        else:
            features = np.asarray(data)
            if features.ndim != 2:
                raise ValueError("feature data must be a two-dimensional array")
            if durations is None or events is None:
                raise ValueError("durations and events are required with array input")
            if len(features) != len(durations) or len(features) != len(events):
                raise ValueError("features, durations, and events must have equal length")
            names = [f"x{index + 1}" for index in range(features.shape[1])]
            frame = pd.DataFrame(features, columns=names)
            frame[duration_col] = np.asarray(durations)
            frame[event_col] = np.asarray(events)

        required = {duration_col, event_col}
        missing = required - set(frame.columns)
        if missing:
            raise ValueError(f"data is missing required columns: {sorted(missing)}")
        rename = {}
        if duration_col != "time":
            rename[duration_col] = "time"
        if event_col != "status":
            rename[event_col] = "status"
        frame = frame.rename(columns=rename)
        duration_col, event_col = "time", "status"

        excluded = {duration_col, event_col}
        if client_col is not None:
            if client_col not in frame:
                raise ValueError(f"client column {client_col!r} was not found")
            excluded.add(client_col)
        feature_names = [name for name in frame.columns if name not in excluded]
        if not feature_names:
            raise ValueError("at least one feature column is required")
        non_numeric = [
            name for name in feature_names if not pd.api.types.is_numeric_dtype(frame[name])
        ]
        if non_numeric:
            raise TypeError(f"feature columns must be numeric: {non_numeric}")
        self.feature_names_ = tuple(feature_names)

        canonical = frame[
            feature_names + [duration_col, event_col] + ([client_col] if client_col else [])
        ]
        if client_col is None:
            return DataSplitter(
                n_clients=self.config.num_clients,
                split_type=self.split_type,
                alpha=self.split_alpha,
                test_size=self.test_size,
                random_state=self.config.random_seed,
            ).split(canonical)
        return self._split_existing_clients(canonical, feature_names, client_col)

    def _split_existing_clients(self, frame, feature_names, client_col) -> DataSet:
        train, test = train_test_split(
            frame,
            test_size=self.test_size,
            random_state=self.config.random_seed,
            stratify=frame["status"],
        )
        clients = {}
        for index, (_, group) in enumerate(train.groupby(client_col, sort=True)):
            x = group[feature_names].to_numpy(dtype=np.float32)
            y = group[["time", "status"]].to_numpy(dtype=np.float32)
            if not len(x):
                continue
            if y[:, 1].sum() < 1:
                raise ValueError(f"client {index} has no observed event after the train/test split")
            clients[f"client{index}"] = (x, y)
        if not clients:
            raise ValueError("no client training data remain after splitting")
        return DataSet(
            clients_set=clients,
            train_data=train[feature_names].to_numpy(dtype=np.float32),
            train_label=train[["time", "status"]].to_numpy(dtype=np.float32),
            test_data=test[feature_names].to_numpy(dtype=np.float32),
            test_label=test[["time", "status"]].to_numpy(dtype=np.float32),
            raw_aug_clients_set={},
        )

    def _coerce_features(self, features) -> np.ndarray:
        if isinstance(features, pd.DataFrame):
            missing = set(self.feature_names_) - set(features.columns)
            if missing:
                raise ValueError(f"prediction data is missing features: {sorted(missing)}")
            array = features[list(self.feature_names_)].to_numpy(dtype=np.float32)
        else:
            array = np.asarray(features, dtype=np.float32)
        if array.ndim != 2 or array.shape[1] != self.config.n_features:
            raise ValueError(
                f"expected a two-dimensional feature matrix with {self.config.n_features} columns"
            )
        return array

    def _require_fitted(self) -> None:
        if self.runner_ is None:
            raise RuntimeError("fit must be called before prediction or evaluation")

fit(data, durations=None, events=None, *, duration_col='time', event_col='status', client_col=None, augmentation=None)

Fit from a package DataSet, DataFrame, or NumPy arrays.

A DataFrame may optionally contain an institution column. Otherwise the requested synthetic partition strategy is applied. The feature count is always inferred from the supplied data.

Source code in federated_survival/api.py

def fit(
    self,
    data,
    durations=None,
    events=None,
    *,
    duration_col: str = "time",
    event_col: str = "status",
    client_col: Optional[str] = None,
    augmentation: Optional[str] = None,
) -> "FederatedSurvival":
    """Fit from a package ``DataSet``, DataFrame, or NumPy arrays.

    A DataFrame may optionally contain an institution column.  Otherwise
    the requested synthetic partition strategy is applied.  The feature
    count is always inferred from the supplied data.
    """

    dataset = self._prepare_dataset(
        data,
        durations,
        events,
        duration_col=duration_col,
        event_col=event_col,
        client_col=client_col,
    )
    first_x = next(iter(dataset.clients_set.values()))[0]
    self.config.n_features = int(first_x.shape[1])
    self.config.num_clients = len(dataset.clients_set)

    runner = FSARunner(self.config)
    if augmentation is None:
        history = runner.run(dataset, type="raw")
    else:
        normalized = augmentation.upper()
        if normalized not in {"MVAEC", "MVAES"}:
            raise ValueError("augmentation must be None, 'MVAEC', or 'MVAES'")
        history = runner.run(dataset, type="raw_aug", aug_method=normalized)

    self.runner_ = runner
    self.history_ = history
    self.dataset_ = dataset
    return self

predict_survival(features)

Return a time-by-sample survival-probability DataFrame.

Source code in federated_survival/api.py

def predict_survival(self, features) -> pd.DataFrame:
    """Return a time-by-sample survival-probability DataFrame."""
    self._require_fitted()
    array = self._coerce_features(features)
    return self.runner_.predict_survival(array)

evaluate(features, durations, events)

Compute time-dependent concordance and IPCW integrated Brier score.

Source code in federated_survival/api.py

def evaluate(self, features, durations, events) -> Dict[str, Any]:
    """Compute time-dependent concordance and IPCW integrated Brier score."""
    survival = self.predict_survival(features)
    duration_values = np.asarray(durations, dtype=float).reshape(-1)
    event_values = np.asarray(events, dtype=float).reshape(-1)
    if survival.shape[1] != len(duration_values) or len(events) != len(duration_values):
        raise ValueError("features, durations, and events must contain the same samples")
    evaluator = EvalSurv(
        survival,
        duration_values,
        event_values,
        censor_surv="km",
    )
    grid = evaluation_time_grid(
        duration_values,
        survival.index.values,
        self.config.evaluation_quantiles,
    )
    diagnostics = self.runner_.last_prediction_diagnostics_
    return {
        "c_index": float(evaluator.concordance_td()),
        "ibs": float(evaluator.integrated_brier_score(grid)),
        "prediction_diagnostics": asdict(diagnostics),
    }

get_run_summary()

Return serialisable configuration and final metrics.

Source code in federated_survival/api.py

def get_run_summary(self) -> Dict[str, Any]:
    """Return serialisable configuration and final metrics."""
    self._require_fitted()
    final = {}
    for name in ("train_Cindex", "train_IBS", "test_Cindex", "test_IBS"):
        values = self.history_.get(name, [])
        final[name] = float(values[-1]) if values else None
    return {
        "model": self.config.model_type,
        "protocol": self.config.federated_protocol,
        "n_clients": self.config.num_clients,
        "n_features": self.config.n_features,
        "global_rounds": self.config.global_epochs,
        "local_steps": self.config.local_epochs,
        "batch_size": self.config.batch_size,
        "final_metrics": final,
    }

get_privacy_info()

Report configured perturbation settings, not an achieved DP budget.

Source code in federated_survival/api.py

def get_privacy_info(self) -> Dict[str, Any]:
    """Report configured perturbation settings, not an achieved DP budget."""
    self._require_fitted()
    config = self.config
    if not config.use_differential_privacy:
        return {"enabled": False, "formal_accounting_available": False}
    from .core.differential_privacy import describe_noise_knobs

    knobs = describe_noise_knobs(config.dp_mechanism)
    return {
        "enabled": True,
        "mechanism": config.dp_mechanism,
        "configured_epsilon": config.dp_epsilon,
        "configured_delta": config.dp_delta if config.dp_mechanism == "gaussian" else None,
        "configured_sensitivity": config.dp_sensitivity,
        "configured_noise_multiplier": (
            config.dp_noise_multiplier if config.dp_mechanism == "gaussian" else None
        ),
        "configured_clip_norm": (
            config.dp_clip_norm if config.dp_mechanism in ("gaussian", "laplace") else None
        ),
        "formal_accounting_available": False,
        "privacy_scope": "experimental clipped client-update perturbation",
        # The Gaussian scale follows dp_noise_multiplier and the Laplace
        # scale follows dp_epsilon, so the reported epsilon is not the
        # noise control under Gaussian.  Naming the driver and the ignored
        # fields keeps that from being misread.
        "noise_driver": knobs["noise_knob"],
        "inactive_parameters": list(knobs["inert_knobs"]),
    }

Composable experiment API

Run one Center/Federated/Local comparison with a compact public call.

Source code in federated_survival/experiment_api.py

def compare_methods(
    dataset: Union[DataSet, Mapping[int, DataSet]],
    *,
    model: str = "DeepSurv",
    protocol: str = "FedAvg",
    protocol_params: Optional[Mapping[str, object]] = None,
    **training_options,
) -> ExperimentResult:
    """Run one Center/Federated/Local comparison with a compact public call."""
    return compare_experiments(
        dataset,
        models=[model],
        protocols=[protocol],
        protocol_params={protocol: protocol_params or {}},
        **training_options,
    )

Cross selected models and protocols on paired per-seed partitions.

Pass one DataSet for a single seed or a mapping returned by partition_data_many for repeated seeds. Within each seed, Center, Federated, and Local fits use the same partition and test set. Incompatible cells are rejected before training. This API does not write an audit manifest; use YAML when a full machine-readable run record is required.

Source code in federated_survival/experiment_api.py

def compare_experiments(
    dataset: Union[DataSet, Mapping[int, DataSet]],
    *,
    models: Sequence[str],
    protocols: Sequence[str],
    global_rounds: int = 10,
    local_steps: int = 5,
    reference_steps: Optional[int] = None,
    batch_size: int = 32,
    learning_rate: float = 1e-3,
    optimizer: str = "adam",
    hidden_nodes: Optional[Sequence[int]] = None,
    num_durations: int = 25,
    activation: str = "relu",
    dropout: float = 0.1,
    weight_decay: float = 0.0,
    client_fraction: float = 1.0,
    seed: int = 42,
    protocol_params: Optional[Mapping[str, Mapping[str, object]]] = None,
    save_predictions: bool = False,
) -> ExperimentResult:
    """Cross selected models and protocols on paired per-seed partitions.

    Pass one ``DataSet`` for a single seed or a mapping returned by
    ``partition_data_many`` for repeated seeds. Within each seed, Center,
    Federated, and Local fits use the same partition and test set. Incompatible
    cells are rejected before training. This API does not write an audit
    manifest; use YAML when a full machine-readable run record is required.
    """
    if isinstance(dataset, DataSet):
        datasets = {int(seed): dataset}
    elif isinstance(dataset, Mapping) and dataset:
        datasets = {int(key): value for key, value in dataset.items()}
        if len(datasets) != len(dataset):
            raise ValueError("dataset seed keys must be unique integers")
    else:
        raise TypeError("dataset must be a DataSet or a non-empty seed-to-DataSet mapping")
    dimensions = {
        current_seed: _validate_dataset(split) for current_seed, split in datasets.items()
    }
    if len({(value[1], value[2]) for value in dimensions.values()}) != 1:
        raise ValueError("all seeds must have the same feature and client counts")
    if isinstance(models, str) or not models or len(set(models)) != len(models):
        raise ValueError("models must be a non-empty sequence without duplicates")
    if isinstance(protocols, str) or not protocols or len(set(protocols)) != len(protocols):
        raise ValueError("protocols must be a non-empty sequence without duplicates")
    if reference_steps is not None and int(reference_steps) <= 0:
        raise ValueError("reference_steps must be positive")
    parameter_map = protocol_params or {}
    if not isinstance(parameter_map, Mapping):
        raise TypeError("protocol_params must map protocol names to parameter mappings")

    configurations: Dict[Tuple[str, str], FSAConfig] = {}
    seeded_configurations: Dict[Tuple[int, str, str], FSAConfig] = {}
    for current_seed, dimensions_for_seed in dimensions.items():
        n_samples, n_features, n_clients = dimensions_for_seed
        for model in models:
            for requested_protocol in protocols:
                config = _make_config(
                    model=model,
                    protocol=requested_protocol,
                    n_samples=n_samples,
                    n_features=n_features,
                    n_clients=n_clients,
                    seed=current_seed,
                    global_rounds=global_rounds,
                    local_steps=local_steps,
                    batch_size=batch_size,
                    learning_rate=learning_rate,
                    optimizer=optimizer,
                    hidden_nodes=hidden_nodes,
                    num_durations=num_durations,
                    activation=activation,
                    dropout=dropout,
                    weight_decay=weight_decay,
                    client_fraction=client_fraction,
                    protocol_params=parameter_map.get(requested_protocol),
                )
                cell_key = (model, config.federated_protocol)
                seeded_key = (current_seed,) + cell_key
                if seeded_key in seeded_configurations:
                    raise ValueError("duplicate model/protocol cell: %s" % (cell_key,))
                seeded_configurations[seeded_key] = config
                configurations[cell_key] = config

    raw_frames = []
    round_rows = []
    prediction_rows = [] if save_predictions else None
    for (current_seed, model, protocol), config in seeded_configurations.items():
        cell_rounds = []
        cell_predictions = [] if save_predictions else None
        cell = run_paired_baselines(
            config,
            datasets[current_seed],
            current_seed,
            baseline_epochs=reference_steps,
            telemetry_rows=cell_rounds,
            prediction_rows=cell_predictions,
        )
        cell["protocol"] = protocol
        raw_frames.append(cell)
        for row in cell_rounds:
            row["protocol"] = protocol
            round_rows.append(row)
        if cell_predictions is not None:
            for row in cell_predictions:
                row["protocol"] = protocol
                prediction_rows.append(row)

    raw = pd.concat(raw_frames, ignore_index=True)
    metrics = raw.copy()
    metrics["method"] = metrics["method"].map(DISPLAY_METHODS)
    return ExperimentResult(
        dataset=next(iter(datasets.values())),
        datasets=datasets,
        metrics=metrics,
        round_metrics=pd.DataFrame(round_rows),
        predictions=pd.DataFrame(prediction_rows or []),
        configurations=configurations,
        _raw_metrics=raw,
    )

Paired metrics, round history, and selectable plots from one or more seeds.

Source code in federated_survival/experiment_api.py

@dataclass
class ExperimentResult:
    """Paired metrics, round history, and selectable plots from one or more seeds."""

    dataset: DataSet
    datasets: Dict[int, DataSet]
    metrics: pd.DataFrame
    round_metrics: pd.DataFrame
    predictions: pd.DataFrame
    configurations: Dict[Tuple[str, str], FSAConfig]
    _raw_metrics: pd.DataFrame

    def _selected_cell(self, model: Optional[str], protocol: Optional[str]) -> Tuple[str, str]:
        cells = list(self.configurations)
        if model is None and protocol is None and len(cells) == 1:
            return cells[0]
        if model is None or protocol is None:
            raise ValueError("select both model and protocol when plotting multiple cells")
        if (model, protocol) not in self.configurations:
            raise ValueError("no result for model=%r, protocol=%r" % (model, protocol))
        return model, protocol

    def plot(
        self,
        kind: str = "final_metrics",
        *,
        model: Optional[str] = None,
        protocol: Optional[str] = None,
        chart: Optional[str] = None,
        seed: Optional[int] = None,
        x: str = "feature_1",
        y: str = "time",
        output_path: Optional[Union[str, Path]] = None,
        show: bool = True,
    ):
        """Plot client partition, final C-index/IBS, or federated round history.

        For multi-model/protocol results, select one cell with ``model`` and
        ``protocol``. Partition plots default to a client-wise strip; set
        ``chart="scatter"`` for feature axes. Boxplots and violin plots need
        at least two paired seeds.
        """
        if kind == "partition":
            selected_seed = next(iter(self.datasets)) if seed is None else int(seed)
            if selected_seed not in self.datasets:
                raise ValueError("no partition for seed %s" % selected_seed)
            return plot_partition(
                self.datasets[selected_seed],
                x=x,
                y=y,
                chart=chart or "strip",
                output_path=output_path,
                show=show,
            )
        if kind not in ("final_metrics", "round_metrics"):
            raise ValueError("kind must be partition, final_metrics, or round_metrics")
        model, protocol = self._selected_cell(model, protocol)

        import matplotlib.pyplot as plt

        from .experiments.visualization import _draw_final_metric, _draw_round_metric

        fig, axes = plt.subplots(1, 2, figsize=(10.6, 4.3))
        if kind == "round_metrics":
            frame = self.round_metrics[
                self.round_metrics["model"].eq(model) & self.round_metrics["protocol"].eq(protocol)
            ]
            if frame.empty:
                plt.close(fig)
                raise ValueError("round metrics are unavailable for this result")
            for ax, metric in zip(axes, ("c_index", "ibs")):
                _draw_round_metric(ax, frame, metric, "mean_ci95")
        else:
            final_chart = chart or "dotplot"
            if final_chart not in PLOT_TYPES:
                plt.close(fig)
                raise ValueError("chart must be one of %s" % (list(PLOT_TYPES),))
            if final_chart in ("boxplot", "violin") and self._raw_metrics["seed"].nunique() < 2:
                plt.close(fig)
                raise ValueError("%s requires at least two paired seeds" % final_chart)
            frame = self._raw_metrics[
                self._raw_metrics["model"].eq(model) & self._raw_metrics["protocol"].eq(protocol)
            ]
            settings = {
                "type": final_chart,
                "methods": ["Center", "Federated", "Weighted Local"],
                "show_raw_points": True,
                "show_mean": True,
            }
            for ax, metric in zip(axes, ("c_index", "ibs")):
                _draw_final_metric(ax, frame, metric, settings)
        fig.tight_layout()
        if output_path is not None:
            path = Path(output_path)
            path.parent.mkdir(parents=True, exist_ok=True)
            fig.savefig(path, dpi=300, bbox_inches="tight")
        if show:
            plt.show()
        return fig

    def save_figures(
        self,
        output_dir: Union[str, Path],
        *,
        combine: bool = False,
        final_chart: str = "dotplot",
        partition_chart: str = "strip",
        model: Optional[str] = None,
        protocol: Optional[str] = None,
        seed: Optional[int] = None,
        x: str = "feature_1",
        y: str = "time",
        image_format: str = "png",
        overwrite: bool = False,
    ) -> Dict[str, Path]:
        """Save separate plots and optionally a dataset-independent combined figure.

        ``combine=False`` writes partition, round-metric, and final-metric
        figures. ``combine=True`` writes those plus one five-panel overview.
        Existing target files are preserved unless ``overwrite=True``.
        """
        model, protocol = self._selected_cell(model, protocol)
        selected_seed = next(iter(self.datasets)) if seed is None else int(seed)
        if selected_seed not in self.datasets:
            raise ValueError("no partition for seed %s" % selected_seed)
        if partition_chart not in ("strip", "scatter"):
            raise ValueError("partition_chart must be strip or scatter")
        if final_chart not in PLOT_TYPES:
            raise ValueError("final_chart must be one of %s" % (list(PLOT_TYPES),))
        if final_chart in ("boxplot", "violin") and len(self.datasets) < 2:
            raise ValueError("%s requires at least two paired seeds" % final_chart)
        if image_format not in ("png", "pdf", "svg"):
            raise ValueError("image_format must be png, pdf, or svg")
        if self.round_metrics.empty:
            raise ValueError("round metrics are unavailable for this result")
        rounds = self.round_metrics[
            self.round_metrics["model"].eq(model) & self.round_metrics["protocol"].eq(protocol)
        ]
        if rounds.empty:
            raise ValueError("round metrics are unavailable for this result")
        raw = self._raw_metrics[
            self._raw_metrics["model"].eq(model) & self._raw_metrics["protocol"].eq(protocol)
        ]
        directory = Path(output_dir)
        names = {
            "partition": "client_partition.%s" % image_format,
            "round_metrics": "round_metrics.%s" % image_format,
            "final_metrics": "final_metrics_%s.%s" % (final_chart, image_format),
        }
        if combine:
            names["combined"] = "combined_%s.%s" % (final_chart, image_format)
        paths = {key: directory / name for key, name in names.items()}
        existing = [str(path) for path in paths.values() if path.exists()]
        if existing and not overwrite:
            raise FileExistsError("figure files already exist: %s" % existing)
        directory.mkdir(parents=True, exist_ok=True)

        import matplotlib.pyplot as plt

        for key, options in (
            (
                "partition",
                {"chart": partition_chart, "seed": selected_seed, "x": x, "y": y},
            ),
            ("round_metrics", {"model": model, "protocol": protocol}),
            (
                "final_metrics",
                {"chart": final_chart, "model": model, "protocol": protocol},
            ),
        ):
            figure = self.plot(kind=key, output_path=paths[key], show=False, **options)
            plt.close(figure)

        if combine:
            from .experiments.visualization import _draw_final_metric, _draw_round_metric

            figure = plt.figure(figsize=(12.5, 13.5), constrained_layout=True)
            grid = figure.add_gridspec(3, 2, height_ratios=(1.25, 1, 1))
            partition_axis = figure.add_subplot(grid[0, :])
            _draw_partition(
                partition_axis,
                self.datasets[selected_seed],
                x,
                y,
                partition_chart,
            )
            round_axes = (figure.add_subplot(grid[1, 0]), figure.add_subplot(grid[1, 1]))
            final_axes = (figure.add_subplot(grid[2, 0]), figure.add_subplot(grid[2, 1]))
            for axis, metric in zip(round_axes, ("c_index", "ibs")):
                _draw_round_metric(axis, rounds, metric, "mean_ci95")
            settings = {
                "type": final_chart,
                "methods": ["Center", "Federated", "Weighted Local"],
                "show_raw_points": True,
                "show_mean": True,
            }
            for axis, metric in zip(final_axes, ("c_index", "ibs")):
                _draw_final_metric(axis, raw, metric, settings)
            figure.savefig(paths["combined"], dpi=300, bbox_inches="tight")
            plt.close(figure)
        return paths

plot(kind='final_metrics', *, model=None, protocol=None, chart=None, seed=None, x='feature_1', y='time', output_path=None, show=True)

Plot client partition, final C-index/IBS, or federated round history.

For multi-model/protocol results, select one cell with model and protocol. Partition plots default to a client-wise strip; set chart="scatter" for feature axes. Boxplots and violin plots need at least two paired seeds.

Source code in federated_survival/experiment_api.py

def plot(
    self,
    kind: str = "final_metrics",
    *,
    model: Optional[str] = None,
    protocol: Optional[str] = None,
    chart: Optional[str] = None,
    seed: Optional[int] = None,
    x: str = "feature_1",
    y: str = "time",
    output_path: Optional[Union[str, Path]] = None,
    show: bool = True,
):
    """Plot client partition, final C-index/IBS, or federated round history.

    For multi-model/protocol results, select one cell with ``model`` and
    ``protocol``. Partition plots default to a client-wise strip; set
    ``chart="scatter"`` for feature axes. Boxplots and violin plots need
    at least two paired seeds.
    """
    if kind == "partition":
        selected_seed = next(iter(self.datasets)) if seed is None else int(seed)
        if selected_seed not in self.datasets:
            raise ValueError("no partition for seed %s" % selected_seed)
        return plot_partition(
            self.datasets[selected_seed],
            x=x,
            y=y,
            chart=chart or "strip",
            output_path=output_path,
            show=show,
        )
    if kind not in ("final_metrics", "round_metrics"):
        raise ValueError("kind must be partition, final_metrics, or round_metrics")
    model, protocol = self._selected_cell(model, protocol)

    import matplotlib.pyplot as plt

    from .experiments.visualization import _draw_final_metric, _draw_round_metric

    fig, axes = plt.subplots(1, 2, figsize=(10.6, 4.3))
    if kind == "round_metrics":
        frame = self.round_metrics[
            self.round_metrics["model"].eq(model) & self.round_metrics["protocol"].eq(protocol)
        ]
        if frame.empty:
            plt.close(fig)
            raise ValueError("round metrics are unavailable for this result")
        for ax, metric in zip(axes, ("c_index", "ibs")):
            _draw_round_metric(ax, frame, metric, "mean_ci95")
    else:
        final_chart = chart or "dotplot"
        if final_chart not in PLOT_TYPES:
            plt.close(fig)
            raise ValueError("chart must be one of %s" % (list(PLOT_TYPES),))
        if final_chart in ("boxplot", "violin") and self._raw_metrics["seed"].nunique() < 2:
            plt.close(fig)
            raise ValueError("%s requires at least two paired seeds" % final_chart)
        frame = self._raw_metrics[
            self._raw_metrics["model"].eq(model) & self._raw_metrics["protocol"].eq(protocol)
        ]
        settings = {
            "type": final_chart,
            "methods": ["Center", "Federated", "Weighted Local"],
            "show_raw_points": True,
            "show_mean": True,
        }
        for ax, metric in zip(axes, ("c_index", "ibs")):
            _draw_final_metric(ax, frame, metric, settings)
    fig.tight_layout()
    if output_path is not None:
        path = Path(output_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(path, dpi=300, bbox_inches="tight")
    if show:
        plt.show()
    return fig

save_figures(output_dir, *, combine=False, final_chart='dotplot', partition_chart='strip', model=None, protocol=None, seed=None, x='feature_1', y='time', image_format='png', overwrite=False)

Save separate plots and optionally a dataset-independent combined figure.

combine=False writes partition, round-metric, and final-metric figures. combine=True writes those plus one five-panel overview. Existing target files are preserved unless overwrite=True.

Source code in federated_survival/experiment_api.py

def save_figures(
    self,
    output_dir: Union[str, Path],
    *,
    combine: bool = False,
    final_chart: str = "dotplot",
    partition_chart: str = "strip",
    model: Optional[str] = None,
    protocol: Optional[str] = None,
    seed: Optional[int] = None,
    x: str = "feature_1",
    y: str = "time",
    image_format: str = "png",
    overwrite: bool = False,
) -> Dict[str, Path]:
    """Save separate plots and optionally a dataset-independent combined figure.

    ``combine=False`` writes partition, round-metric, and final-metric
    figures. ``combine=True`` writes those plus one five-panel overview.
    Existing target files are preserved unless ``overwrite=True``.
    """
    model, protocol = self._selected_cell(model, protocol)
    selected_seed = next(iter(self.datasets)) if seed is None else int(seed)
    if selected_seed not in self.datasets:
        raise ValueError("no partition for seed %s" % selected_seed)
    if partition_chart not in ("strip", "scatter"):
        raise ValueError("partition_chart must be strip or scatter")
    if final_chart not in PLOT_TYPES:
        raise ValueError("final_chart must be one of %s" % (list(PLOT_TYPES),))
    if final_chart in ("boxplot", "violin") and len(self.datasets) < 2:
        raise ValueError("%s requires at least two paired seeds" % final_chart)
    if image_format not in ("png", "pdf", "svg"):
        raise ValueError("image_format must be png, pdf, or svg")
    if self.round_metrics.empty:
        raise ValueError("round metrics are unavailable for this result")
    rounds = self.round_metrics[
        self.round_metrics["model"].eq(model) & self.round_metrics["protocol"].eq(protocol)
    ]
    if rounds.empty:
        raise ValueError("round metrics are unavailable for this result")
    raw = self._raw_metrics[
        self._raw_metrics["model"].eq(model) & self._raw_metrics["protocol"].eq(protocol)
    ]
    directory = Path(output_dir)
    names = {
        "partition": "client_partition.%s" % image_format,
        "round_metrics": "round_metrics.%s" % image_format,
        "final_metrics": "final_metrics_%s.%s" % (final_chart, image_format),
    }
    if combine:
        names["combined"] = "combined_%s.%s" % (final_chart, image_format)
    paths = {key: directory / name for key, name in names.items()}
    existing = [str(path) for path in paths.values() if path.exists()]
    if existing and not overwrite:
        raise FileExistsError("figure files already exist: %s" % existing)
    directory.mkdir(parents=True, exist_ok=True)

    import matplotlib.pyplot as plt

    for key, options in (
        (
            "partition",
            {"chart": partition_chart, "seed": selected_seed, "x": x, "y": y},
        ),
        ("round_metrics", {"model": model, "protocol": protocol}),
        (
            "final_metrics",
            {"chart": final_chart, "model": model, "protocol": protocol},
        ),
    ):
        figure = self.plot(kind=key, output_path=paths[key], show=False, **options)
        plt.close(figure)

    if combine:
        from .experiments.visualization import _draw_final_metric, _draw_round_metric

        figure = plt.figure(figsize=(12.5, 13.5), constrained_layout=True)
        grid = figure.add_gridspec(3, 2, height_ratios=(1.25, 1, 1))
        partition_axis = figure.add_subplot(grid[0, :])
        _draw_partition(
            partition_axis,
            self.datasets[selected_seed],
            x,
            y,
            partition_chart,
        )
        round_axes = (figure.add_subplot(grid[1, 0]), figure.add_subplot(grid[1, 1]))
        final_axes = (figure.add_subplot(grid[2, 0]), figure.add_subplot(grid[2, 1]))
        for axis, metric in zip(round_axes, ("c_index", "ibs")):
            _draw_round_metric(axis, rounds, metric, "mean_ci95")
        settings = {
            "type": final_chart,
            "methods": ["Center", "Federated", "Weighted Local"],
            "show_raw_points": True,
            "show_mean": True,
        }
        for axis, metric in zip(final_axes, ("c_index", "ibs")):
            _draw_final_metric(axis, raw, metric, settings)
        figure.savefig(paths["combined"], dpi=300, bbox_inches="tight")
        plt.close(figure)
    return paths

Interactive data functions

Generate one simulated survival table through the public package API.

censoring is the censoring-distribution control used by Weibull, log-normal, and SDGM1 simulations; it is not an exact requested censoring proportion. SDGM2/SDGM3 use u_max and SDGM4 uses c_step.

Source code in federated_survival/data/workflow.py

def simulate_data(
    mechanism: str = "weibull",
    n_samples: int = 300,
    n_features: int = 5,
    censoring: float = 0.4,
    seed: Optional[int] = 42,
    *,
    u_max: float = 4.0,
    c_step: float = 0.4,
) -> pd.DataFrame:
    """Generate one simulated survival table through the public package API.

    ``censoring`` is the censoring-distribution control used by Weibull,
    log-normal, and SDGM1 simulations; it is not an exact requested censoring
    proportion. SDGM2/SDGM3 use ``u_max`` and SDGM4 uses ``c_step``.
    """
    if int(n_samples) <= 0 or int(n_features) <= 0:
        raise ValueError("n_samples and n_features must be positive")
    generator = DataGenerator(
        SimulationConfig(
            n_samples=int(n_samples),
            n_features=int(n_features),
            random_state=seed,
        )
    )
    return generator.generate(
        mechanism,
        c_mean=float(censoring),
        u_max=float(u_max),
        c_step=float(c_step),
    )

Load CSV/Excel survival data and return canonical x1..xp/time/status columns.

Source code in federated_survival/data/workflow.py

def load_data(
    path: Union[str, Path],
    duration_column: str = "time",
    event_column: str = "status",
) -> pd.DataFrame:
    """Load CSV/Excel survival data and return canonical x1..xp/time/status columns."""
    if duration_column == event_column:
        raise ValueError("duration_column and event_column must be different")
    frame = DataLoader(
        time_column=duration_column,
        status_column=event_column,
    ).load(path)
    return frame.rename(columns={duration_column: "time", event_column: "status"})

Load a real survival dataset bundled with the installed package.

Bundled tables live inside the package itself (federated_survival/data/real/), so they are available after a normal pip install -- no repository clone or manual data download needed. Use :func:available_real_datasets to list the shipped names.

Parameters:

Name	Type	Description	Default
name	str	Dataset name, e.g. "gbsg" for the GBSG breast-cancer table.	'gbsg'
duration_column	str	Column treated as survival time (kept as "time").	'time'
event_column	str	Column treated as the event indicator (kept as "status").	'status'

Returns:

Name	Type	Description
	DataFrame	pd.DataFrame: Canonical x1..xp/time/status table, identical to
what	DataFrame	func:load_data returns.

Source code in federated_survival/data/workflow.py

def load_real_data(
    name: str = "gbsg",
    *,
    duration_column: str = "time",
    event_column: str = "status",
) -> pd.DataFrame:
    """Load a real survival dataset bundled with the installed package.

    Bundled tables live inside the package itself
    (``federated_survival/data/real/``), so they are available after a normal
    ``pip install`` -- no repository clone or manual data download needed.
    Use :func:`available_real_datasets` to list the shipped names.

    Args:
        name: Dataset name, e.g. ``"gbsg"`` for the GBSG breast-cancer table.
        duration_column: Column treated as survival time (kept as ``"time"``).
        event_column: Column treated as the event indicator (kept as ``"status"``).

    Returns:
        pd.DataFrame: Canonical ``x1..xp/time/status`` table, identical to
        what :func:`load_data` returns.
    """
    if not name or Path(name).name != name:
        raise ValueError(f"Invalid dataset name: {name!r}")
    path = _BUNDLED_REAL_DATA_DIR / f"{name}.csv"
    if not path.exists():
        available = ", ".join(available_real_datasets()) or "(none)"
        raise FileNotFoundError(
            f"Dataset {name!r} is not bundled with the package. Available: {available}"
        )
    return load_data(path, duration_column=duration_column, event_column=event_column)

Return the sorted names of real datasets bundled with the package.

Source code in federated_survival/data/workflow.py

def available_real_datasets() -> list:
    """Return the sorted names of real datasets bundled with the package."""
    if not _BUNDLED_REAL_DATA_DIR.is_dir():
        return []
    return sorted(p.stem for p in _BUNDLED_REAL_DATA_DIR.glob("*.csv"))

Split data and optionally preprocess features from training statistics only.

Source code in federated_survival/data/workflow.py

def partition_data(
    data: pd.DataFrame,
    n_clients: int = 3,
    method: str = "iid",
    test_size: float = 0.2,
    seed: Optional[int] = 42,
    alpha: float = 0.5,
    *,
    standardize: bool = False,
    missing_values: str = "error",
) -> DataSet:
    """Split data and optionally preprocess features from training statistics only."""
    if missing_values not in ("error", "drop", "median"):
        raise ValueError("missing_values must be error, drop, or median")
    feature_columns = [column for column in data if column not in ("time", "status")]
    if missing_values == "error" and data[feature_columns].isna().any().any():
        raise ValueError("feature columns contain missing values; choose drop or median")
    if missing_values == "drop":
        data = data.dropna(axis=0).reset_index(drop=True)
    dataset = DataSplitter(
        n_clients=int(n_clients),
        split_type=method,
        alpha=float(alpha),
        test_size=float(test_size),
        random_state=seed,
    ).split(data)
    return preprocess_partition(dataset, missing_values=missing_values, standardize=standardize)

Create paired client partitions for multiple random seeds.

Source code in federated_survival/data/workflow.py

def partition_data_many(
    data: pd.DataFrame,
    seeds: Sequence[int],
    n_clients: int = 3,
    method: str = "iid",
    test_size: float = 0.2,
    alpha: float = 0.5,
    *,
    standardize: bool = False,
    missing_values: str = "error",
) -> Dict[int, DataSet]:
    """Create paired client partitions for multiple random seeds."""
    if isinstance(seeds, (str, bytes)) or not seeds:
        raise ValueError("seeds must be a non-empty sequence of integers")
    parsed = [int(seed) for seed in seeds]
    if len(set(parsed)) != len(parsed):
        raise ValueError("seeds must not contain duplicates")
    return {
        seed: partition_data(
            data,
            n_clients=n_clients,
            method=method,
            test_size=test_size,
            seed=seed,
            alpha=alpha,
            standardize=standardize,
            missing_values=missing_values,
        )
        for seed in parsed
    }

Draw a feature scatter or a client-wise event/censoring strip plot.

The returned object is a Matplotlib Figure. Set show=False for scripts, tests, or headless execution and use output_path to save it.

Source code in federated_survival/data/workflow.py

def plot_partition(
    dataset: DataSet,
    x: str = "feature_1",
    y: str = "time",
    output_path: Optional[Union[str, Path]] = None,
    show: bool = True,
    *,
    chart: str = "scatter",
):
    """Draw a feature scatter or a client-wise event/censoring strip plot.

    The returned object is a Matplotlib ``Figure``. Set ``show=False`` for
    scripts, tests, or headless execution and use ``output_path`` to save it.
    """
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(figsize=(8.4, 5.2))
    _draw_partition(ax, dataset, x, y, chart)
    fig.tight_layout()

    if output_path is not None:
        path = Path(output_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(path, dpi=300, bbox_inches="tight")
    if show:
        plt.show()
    return fig

Return original-plus-synthetic client arrays without altering dataset.

MVAEC augments each client locally. MVAES pools generated samples before redistributing them; it must not be presented as a no-sharing workflow. The existing MVAE trainer uses 500 epochs per client. The output is for exploratory experiments, not a privacy guarantee. By default, synthetic samples are conditionally drawn near latent codes from the wider half of each client's uncensored-time distribution. Use sampling="unconditional" to recover prior-only latent sampling.

Source code in federated_survival/data/augmentation_workflow.py

def augment_clients(
    dataset: DataSet,
    method: str = "MVAEC",
    *,
    k: float = 0.5,
    latent_num: int = 10,
    hidden_num: int = 30,
    alpha: float = 1.0,
    beta: float = 1.0,
    sampling: str = "sparse",
    sparse_gamma: float = 0.1,
    seed: Optional[int] = 42,
) -> ClientArrays:
    """Return original-plus-synthetic client arrays without altering ``dataset``.

    MVAEC augments each client locally. MVAES pools generated samples before
    redistributing them; it must not be presented as a no-sharing workflow.
    The existing MVAE trainer uses 500 epochs per client. The output is for
    exploratory experiments, not a privacy guarantee. By default, synthetic
    samples are conditionally drawn near latent codes from the wider half of
    each client's uncensored-time distribution. Use
    ``sampling="unconditional"`` to recover prior-only latent sampling.
    """
    if not isinstance(dataset, DataSet):
        raise TypeError("dataset must be the result of partition_data(...)")
    selected = str(method).upper()
    if selected not in ("MVAEC", "MVAES"):
        raise ValueError("method must be MVAEC or MVAES")
    if not 0 < float(k) <= 1:
        raise ValueError("k must be in (0, 1]")
    if latent_num <= 0 or hidden_num <= 0 or alpha <= 0 or beta <= 0:
        raise ValueError("latent_num, hidden_num, alpha, and beta must be positive")
    if str(sampling).lower().replace("-", "_") not in {"sparse", "unconditional"}:
        raise ValueError("sampling must be 'sparse' or 'unconditional'")
    if sparse_gamma < 0:
        raise ValueError("sparse_gamma must be non-negative")

    augmenter = DataAugmenter(
        latent_num=int(latent_num),
        hidden_num=int(hidden_num),
        alpha=float(alpha),
        beta=float(beta),
        sampling=sampling,
        sparse_gamma=float(sparse_gamma),
    )
    if seed is None:
        return getattr(augmenter, selected.lower())(dataset.clients_set, k=float(k))

    # The legacy MVAE and redistribution code uses process-wide RNGs. Keep a
    # reproducible local call from changing the caller's random streams.
    python_state = random.getstate()
    numpy_state = np.random.get_state()
    torch_state = torch.random.get_rng_state()
    try:
        random.seed(int(seed))
        np.random.seed(int(seed))
        torch.manual_seed(int(seed))
        return getattr(augmenter, selected.lower())(dataset.clients_set, k=float(k))
    finally:
        random.setstate(python_state)
        np.random.set_state(numpy_state)
        torch.random.set_rng_state(torch_state)

Compare generated samples from augmentation methods in separate panels.

chart="strip" places clients on the x-axis and time on the y-axis, distinguishing original events, original censoring, and synthetic events. chart="scatter" uses feature/time axes and colours for clients. The held-out test set is never plotted or augmented.

Source code in federated_survival/data/augmentation_workflow.py

def plot_augmentation_comparison(
    dataset: DataSet,
    augmentations: Mapping[str, ClientArrays],
    *,
    chart: str = "scatter",
    x: Optional[str] = None,
    y: str = "time",
    output_path: Optional[Union[str, Path]] = None,
    show: bool = True,
):
    """Compare generated samples from augmentation methods in separate panels.

    ``chart="strip"`` places clients on the x-axis and time on the y-axis,
    distinguishing original events, original censoring, and synthetic events.
    ``chart="scatter"`` uses feature/time axes and colours for clients.
    The held-out test set is never plotted or augmented.
    """
    if not isinstance(dataset, DataSet):
        raise TypeError("dataset must be the result of partition_data(...)")
    if not isinstance(augmentations, Mapping) or not augmentations:
        raise ValueError("augmentations must map method names to augmented client arrays")
    if chart not in ("scatter", "strip"):
        raise ValueError("chart must be scatter or strip")
    if chart == "strip" and (x is not None or y != "time"):
        raise ValueError("strip chart fixes x=client and y=time; omit x and y")
    scatter_x = x or "feature_1"
    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D

    figure, axes = plt.subplots(
        1,
        len(augmentations),
        figsize=(6.2 * len(augmentations), 5.2),
        sharex=True,
        sharey=True,
        squeeze=False,
    )
    palette = plt.get_cmap("tab10")
    try:
        for ax, (method, clients) in zip(axes[0], augmentations.items()):
            if set(clients) != set(dataset.clients_set):
                raise ValueError("augmented clients must match the original client IDs")
            synthetic_total = 0
            rng = np.random.RandomState(20260922)
            for index, (client_id, (raw_x, raw_y)) in enumerate(dataset.clients_set.items()):
                augmented_x, augmented_y = clients[client_id]
                raw_x, raw_y = np.asarray(raw_x), np.asarray(raw_y)
                augmented_x, augmented_y = np.asarray(augmented_x), np.asarray(augmented_y)
                n_raw = len(raw_x)
                if (
                    augmented_x.ndim != 2
                    or augmented_y.ndim != 2
                    or augmented_x.shape[1] != raw_x.shape[1]
                    or augmented_y.shape[1] != 2
                    or len(augmented_x) != len(augmented_y)
                    or len(augmented_x) < n_raw
                    or not np.array_equal(augmented_x[:n_raw], raw_x)
                    or not np.array_equal(augmented_y[:n_raw], raw_y)
                ):
                    raise ValueError("augmented arrays must start with unchanged original samples")
                synthetic_x, synthetic_y = augmented_x[n_raw:], augmented_y[n_raw:]
                synthetic_total += len(synthetic_x)
                if chart == "strip":
                    position = index + 1
                    for status, colour, marker in (
                        (1, "#0072B2", "o"),
                        (0, "#D55E00", "x"),
                    ):
                        selected = raw_y[:, 1].astype(int) == status
                        ax.scatter(
                            position + rng.uniform(-0.13, 0.13, int(selected.sum())),
                            raw_y[selected, 0],
                            color=colour,
                            marker=marker,
                            s=20,
                            alpha=0.58,
                            linewidths=0.8,
                        )
                    if len(synthetic_y):
                        ax.scatter(
                            position + rng.uniform(-0.13, 0.13, len(synthetic_y)),
                            synthetic_y[:, 0],
                            color="#009E73",
                            marker="^",
                            s=32,
                            alpha=0.85,
                        )
                else:
                    colour = palette(index)
                    ax.scatter(
                        _axis_values(raw_x, raw_y, scatter_x),
                        _axis_values(raw_x, raw_y, y),
                        s=13,
                        alpha=0.24,
                        color=colour,
                        marker="o",
                    )
                    if len(synthetic_x):
                        ax.scatter(
                            _axis_values(synthetic_x, synthetic_y, scatter_x),
                            _axis_values(synthetic_x, synthetic_y, y),
                            s=38,
                            alpha=0.85,
                            color=colour,
                            marker="x",
                            linewidths=1,
                        )
            ax.set_title("%s (%d synthetic)" % (method, synthetic_total), loc="left")
            if chart == "strip":
                ax.set_xticks(
                    range(1, len(dataset.clients_set) + 1),
                    ["Client %d" % (index + 1) for index in range(len(dataset.clients_set))],
                )
                ax.set_xlim(0.5, len(dataset.clients_set) + 0.5)
                ax.set_xlabel("Client")
            else:
                ax.set_xlabel("Time" if scatter_x == "time" else scatter_x.replace("_", " "))
            ax.grid(alpha=0.2)
        axes[0, 0].set_ylabel("Time" if y == "time" else str(y).replace("_", " "))
        if chart == "strip":
            handles = [
                Line2D(
                    [],
                    [],
                    marker="o",
                    linestyle="None",
                    color="#0072B2",
                    label="Original, uncensored",
                ),
                Line2D(
                    [],
                    [],
                    marker="x",
                    linestyle="None",
                    color="#D55E00",
                    label="Original, censored",
                ),
                Line2D(
                    [],
                    [],
                    marker="^",
                    linestyle="None",
                    color="#009E73",
                    label="Synthetic, uncensored",
                ),
            ]
        else:
            handles = [
                Line2D(
                    [], [], marker="o", linestyle="None", color="gray", alpha=0.5, label="Original"
                ),
                Line2D([], [], marker="x", linestyle="None", color="gray", label="Synthetic"),
            ]
            handles.extend(
                Line2D(
                    [],
                    [],
                    marker="o",
                    linestyle="None",
                    color=palette(index),
                    label="Client %d" % (index + 1),
                )
                for index in range(len(dataset.clients_set))
            )
        figure.legend(handles=handles, loc="upper center", ncol=len(handles), frameon=True)
        figure.tight_layout(rect=(0, 0, 1, 0.92))
        if output_path is not None:
            path = Path(output_path)
            path.parent.mkdir(parents=True, exist_ok=True)
            figure.savefig(path, dpi=300, bbox_inches="tight")
        if show:
            plt.show()
        return figure
    except Exception:
        plt.close(figure)
        raise

Model adapters

Bases: ABC

Contract implemented by every built-in and third-party survival model.

Source code in federated_survival/models/adapters.py

class ModelAdapter(ABC):
    """Contract implemented by every built-in and third-party survival model."""
    name: str
    requires_baseline_hazards: bool = False
    uses_label_transform: bool = False

    def create_label_transform(self, config):
        return None

    def configure_targets(self, config, durations, events):
        """Fit any label transform and update output dimensionality."""
        transform = self.create_label_transform(config)
        if transform is None:
            config.labtrans = None
            config.out_features = 1
            return None
        transform.fit_transform(durations, events)
        config.labtrans = transform
        config.out_features = transform.out_features
        return transform

    def transform_target(self, labels: np.ndarray, label_transform=None):
        durations = np.asarray(labels)[:, 0]
        events = np.asarray(labels)[:, 1]
        if label_transform is None:
            return durations, events
        return label_transform.transform(durations, events)

    def build_network(self, config) -> torch.nn.Module:
        return tt.practical.MLPVanilla(
            in_features=config.n_features,
            num_nodes=config.num_nodes,
            out_features=config.out_features,
            batch_norm=config.batch_norm,
            dropout=config.dropout,
            activation=config.activation,
            output_bias=False,
        )

    @abstractmethod
    def build_model(self, network, config, label_transform=None, optimizer=None):
        """Wrap a Torch network in the corresponding PyCox model."""

    def prepare_for_prediction(
        self,
        model,
        train_x: np.ndarray,
        train_y: np.ndarray,
        label_transform=None,
    ) -> None:
        """Attach training data and estimate a Cox baseline when required."""
        if not self.requires_baseline_hazards:
            return
        target = self.transform_target(train_y, label_transform)
        model.fit(
            train_x,
            target,
            batch_size=len(train_x),
            epochs=0,
            verbose=False,
        )
        model.compute_baseline_hazards()

    def predict_survival(self, model, features, *, validation_mode="raise"):
        survival = model.predict_surv_df(np.asarray(features, dtype=np.float32))
        diagnostics = validate_survival_predictions(survival, mode=validation_mode)
        return survival, diagnostics

configure_targets(config, durations, events)

Fit any label transform and update output dimensionality.

Source code in federated_survival/models/adapters.py

def configure_targets(self, config, durations, events):
    """Fit any label transform and update output dimensionality."""
    transform = self.create_label_transform(config)
    if transform is None:
        config.labtrans = None
        config.out_features = 1
        return None
    transform.fit_transform(durations, events)
    config.labtrans = transform
    config.out_features = transform.out_features
    return transform

build_model(network, config, label_transform=None, optimizer=None) abstractmethod

Wrap a Torch network in the corresponding PyCox model.

Source code in federated_survival/models/adapters.py

@abstractmethod
def build_model(self, network, config, label_transform=None, optimizer=None):
    """Wrap a Torch network in the corresponding PyCox model."""

prepare_for_prediction(model, train_x, train_y, label_transform=None)

Attach training data and estimate a Cox baseline when required.

Source code in federated_survival/models/adapters.py

def prepare_for_prediction(
    self,
    model,
    train_x: np.ndarray,
    train_y: np.ndarray,
    label_transform=None,
) -> None:
    """Attach training data and estimate a Cox baseline when required."""
    if not self.requires_baseline_hazards:
        return
    target = self.transform_target(train_y, label_transform)
    model.fit(
        train_x,
        target,
        batch_size=len(train_x),
        epochs=0,
        verbose=False,
    )
    model.compute_baseline_hazards()

Register a custom adapter under a case-insensitive name.

Source code in federated_survival/models/adapters.py

def register_model_adapter(
    name: str,
    adapter: Type[ModelAdapter],
    *,
    replace: bool = False,
) -> None:
    """Register a custom adapter under a case-insensitive name."""
    if not isinstance(adapter, type) or not issubclass(adapter, ModelAdapter):
        raise TypeError("adapter must be a ModelAdapter subclass")
    normalized = name.strip().lower()
    if not normalized:
        raise ValueError("adapter name cannot be empty")
    if normalized in _REGISTRY and not replace:
        raise ValueError(f"an adapter named {name!r} is already registered")
    _REGISTRY[normalized] = adapter
    _CANONICAL_NAMES[normalized] = name

Validate probability bounds, finiteness, and monotonicity.

Identical columns are recorded as a degeneracy diagnostic rather than an unconditional error because a legitimately uninformative model may emit the same curve for every sample. The experiment layer can nevertheless fail a run when that behaviour is unexpected.

Source code in federated_survival/models/adapters.py

def validate_survival_predictions(
    survival: pd.DataFrame,
    *,
    mode: str = "raise",
    tolerance: float = 1e-6,
) -> PredictionDiagnostics:
    """Validate probability bounds, finiteness, and monotonicity.

    Identical columns are recorded as a degeneracy diagnostic rather than an
    unconditional error because a legitimately uninformative model may emit
    the same curve for every sample.  The experiment layer can nevertheless
    fail a run when that behaviour is unexpected.
    """

    if mode not in {"raise", "warn", "none"}:
        raise ValueError("prediction validation mode must be 'raise', 'warn', or 'none'")
    if not isinstance(survival, pd.DataFrame):
        raise TypeError("survival predictions must be a pandas DataFrame")

    values = survival.to_numpy(dtype=float)
    finite = bool(values.size and np.isfinite(values).all())
    within_bounds = bool(
        finite
        and np.all(values >= -tolerance)
        and np.all(values <= 1.0 + tolerance)
    )
    monotone = bool(
        finite
        and values.shape[0] >= 2
        and np.all(np.diff(values, axis=0) <= tolerance)
    )
    identical = bool(
        values.shape[1] > 1
        and finite
        and np.allclose(values, values[:, [0]], atol=tolerance, rtol=0.0)
    )
    diagnostics = PredictionDiagnostics(
        n_times=int(values.shape[0]) if values.ndim == 2 else 0,
        n_samples=int(values.shape[1]) if values.ndim == 2 else 0,
        finite=finite,
        within_probability_bounds=within_bounds,
        monotone_non_increasing=monotone,
        identical_across_samples=identical,
    )
    if not diagnostics.valid and mode != "none":
        message = (
            "Invalid survival predictions: "
            f"shape={values.shape}, finite={finite}, "
            f"within_probability_bounds={within_bounds}, "
            f"monotone_non_increasing={monotone}. "
            "Inspect the learning rate, privacy-noise scale, and model-specific target settings."
        )
        if mode == "raise":
            raise PredictionValidationError(message)
        warnings.warn(message, RuntimeWarning, stacklevel=2)
    return diagnostics

Federated protocols

Bases: ABC

Base class for built-in and third-party federated protocols.

Source code in federated_survival/protocols/base.py

class FederatedProtocol(ABC):
    """Base class for built-in and third-party federated protocols."""

    name: str
    capabilities = ProtocolCapabilities()
    owns_training_loop: bool = False

    def __init__(self) -> None:
        self.config = None

    def validate_config(self, config) -> None:
        supported = self.capabilities.supported_models
        if supported is not None and config.model_type not in supported:
            choices = ", ".join(supported)
            raise ValueError(
                f"protocol {self.name!r} supports only these models: {choices}"
            )
        if (
            not self.capabilities.supports_partial_participation
            and config.client_sample_ratio != 1.0
        ):
            raise ValueError(f"protocol {self.name!r} requires full client participation")
        if not self.capabilities.supports_minibatch and not config.full_batch:
            # A protocol-owned loop may consume complete local risk-set
            # summaries without using the generic ``full_batch`` switch.
            if not self.owns_training_loop:
                raise ValueError(f"protocol {self.name!r} does not support mini-batches")

    def begin_run(self, config, global_state: Mapping[str, torch.Tensor]) -> None:
        """Reset protocol state before a new fit."""
        self.config = config

    def prepare_client_model(self, model, global_state: Mapping[str, torch.Tensor]) -> None:
        """Optionally modify the local objective before optimizer steps."""

    def aggregate(
        self,
        global_state: Mapping[str, torch.Tensor],
        client_states: Sequence[Mapping[str, torch.Tensor]],
        client_weights: Sequence[float],
        *,
        round_index: int,
    ) -> Dict[str, torch.Tensor]:
        raise NotImplementedError

    def round_communication_bytes(self, n_clients: int, model_bytes: int) -> int:
        """Return the payload-byte estimate for one parameter-exchange round."""
        return int(2 * n_clients * model_bytes)

    def metadata(self) -> Dict[str, object]:
        return {
            "name": self.name,
            "capabilities": asdict(self.capabilities),
        }

begin_run(config, global_state)

Reset protocol state before a new fit.

Source code in federated_survival/protocols/base.py

def begin_run(self, config, global_state: Mapping[str, torch.Tensor]) -> None:
    """Reset protocol state before a new fit."""
    self.config = config

prepare_client_model(model, global_state)

Optionally modify the local objective before optimizer steps.

Source code in federated_survival/protocols/base.py

def prepare_client_model(self, model, global_state: Mapping[str, torch.Tensor]) -> None:
    """Optionally modify the local objective before optimizer steps."""

round_communication_bytes(n_clients, model_bytes)

Return the payload-byte estimate for one parameter-exchange round.

Source code in federated_survival/protocols/base.py

def round_communication_bytes(self, n_clients: int, model_bytes: int) -> int:
    """Return the payload-byte estimate for one parameter-exchange round."""
    return int(2 * n_clients * model_bytes)

Machine-readable compatibility declaration for a protocol.

Source code in federated_survival/protocols/base.py

@dataclass(frozen=True)
class ProtocolCapabilities:
    """Machine-readable compatibility declaration for a protocol."""

    partition_geometry: str = "horizontal"
    supported_models: Optional[Tuple[str, ...]] = None
    supports_minibatch: bool = True
    supports_partial_participation: bool = True
    exchanges_model_parameters: bool = True
    requires_risk_set_statistics: bool = False

Register a protocol class under a case-insensitive public name.

Source code in federated_survival/protocols/base.py

def register_federated_protocol(
    name: str,
    protocol: ProtocolType,
    *,
    aliases: Sequence[str] = (),
    replace: bool = False,
) -> None:
    """Register a protocol class under a case-insensitive public name."""
    if not isinstance(protocol, type) or not issubclass(protocol, FederatedProtocol):
        raise TypeError("protocol must be a FederatedProtocol subclass")
    normalized = str(name).strip().lower()
    if not normalized:
        raise ValueError("protocol name cannot be empty")
    if normalized in _REGISTRY and not replace:
        raise ValueError(f"a protocol named {name!r} is already registered")
    _REGISTRY[normalized] = protocol
    _CANONICAL_NAMES[normalized] = name
    for alias in aliases:
        alias_key = str(alias).strip().lower()
        if not alias_key:
            raise ValueError("protocol alias cannot be empty")
        if alias_key in _ALIASES and not replace:
            raise ValueError(f"a protocol alias named {alias!r} is already registered")
        _ALIASES[alias_key] = normalized

Bases: FederatedProtocol

Distributed Newton solver for a linear pooled CoxPH objective.

Source code in federated_survival/protocols/webdisco.py

class WebDISCOStyleCoxProtocol(FederatedProtocol):
    """Distributed Newton solver for a linear pooled CoxPH objective."""

    name = "WebDISCO-style"
    owns_training_loop = True
    capabilities = ProtocolCapabilities(
        partition_geometry="horizontal",
        supported_models=("CoxPH",),
        supports_minibatch=False,
        supports_partial_participation=False,
        exchanges_model_parameters=False,
        requires_risk_set_statistics=True,
    )

    def validate_config(self, config) -> None:
        super().validate_config(config)
        if config.use_differential_privacy:
            raise ValueError(
                "WebDISCO-style does not use the client-update perturbation module"
            )

    @staticmethod
    def _coerce_clients(client_sets: Mapping) -> Dict[str, Tuple[np.ndarray, np.ndarray]]:
        result = {}
        feature_count = None
        for client_id, (features, labels) in client_sets.items():
            x = np.asarray(features, dtype=np.float64)
            y = np.asarray(labels, dtype=np.float64)
            if x.ndim != 2 or y.ndim != 2 or y.shape[1] != 2 or len(x) != len(y):
                raise ValueError("each WebDISCO client must contain X and [time, event] arrays")
            if not np.isfinite(x).all() or not np.isfinite(y).all():
                raise ValueError("WebDISCO-style inputs must be finite")
            if not np.isin(y[:, 1], [0.0, 1.0]).all():
                raise ValueError("event indicators must be binary")
            if feature_count is None:
                feature_count = x.shape[1]
            elif x.shape[1] != feature_count:
                raise ValueError("all WebDISCO clients must have the same feature dimension")
            result[str(client_id)] = (x, y)
        if not result:
            raise ValueError("WebDISCO-style requires at least one client")
        if sum(int(y[:, 1].sum()) for _, y in result.values()) == 0:
            raise ValueError("WebDISCO-style requires at least one observed event")
        return result

    @staticmethod
    def event_times(client_sets: Mapping) -> np.ndarray:
        clients = WebDISCOStyleCoxProtocol._coerce_clients(client_sets)
        values = [y[y[:, 1] > 0, 0] for _, y in clients.values()]
        return np.unique(np.concatenate(values))

    @staticmethod
    def _risk_maxima(x, y, beta, event_times):
        linear_predictor = x @ beta
        maxima = np.full(len(event_times), -np.inf, dtype=np.float64)
        durations = y[:, 0]
        for index, time in enumerate(event_times):
            risk = durations >= time
            if risk.any():
                maxima[index] = float(np.max(linear_predictor[risk]))
        return maxima

    @staticmethod
    def _local_statistics(x, y, beta, event_times, global_maxima):
        n_times = len(event_times)
        n_features = x.shape[1]
        s0 = np.zeros(n_times, dtype=np.float64)
        s1 = np.zeros((n_times, n_features), dtype=np.float64)
        s2 = np.zeros((n_times, n_features, n_features), dtype=np.float64)
        event_count = np.zeros(n_times, dtype=np.float64)
        event_x_sum = np.zeros((n_times, n_features), dtype=np.float64)
        event_eta_sum = np.zeros(n_times, dtype=np.float64)
        durations, events = y[:, 0], y[:, 1]
        eta = x @ beta
        for index, time in enumerate(event_times):
            risk = durations >= time
            event = (durations == time) & (events > 0)
            shifted = np.exp(eta[risk] - global_maxima[index])
            risk_x = x[risk]
            s0[index] = shifted.sum()
            s1[index] = np.sum(risk_x * shifted[:, None], axis=0)
            s2[index] = np.einsum("ni,nj,n->ij", risk_x, risk_x, shifted)
            event_count[index] = event.sum()
            if event.any():
                event_x_sum[index] = x[event].sum(axis=0)
                event_eta_sum[index] = eta[event].sum()
        return s0, s1, s2, event_count, event_x_sum, event_eta_sum

    def distributed_objective(self, client_sets, beta, event_times=None):
        """Return pooled Breslow log likelihood, score, information and bytes."""
        clients = self._coerce_clients(client_sets)
        beta = np.asarray(beta, dtype=np.float64).reshape(-1)
        if event_times is None:
            event_times = self.event_times(clients)
        event_times = np.asarray(event_times, dtype=np.float64)
        local_maxima = np.stack(
            [self._risk_maxima(x, y, beta, event_times) for x, y in clients.values()]
        )
        global_maxima = np.max(local_maxima, axis=0)
        if not np.isfinite(global_maxima).all():
            raise FloatingPointError("an event time has an empty global risk set")

        totals = None
        for x, y in clients.values():
            values = self._local_statistics(x, y, beta, event_times, global_maxima)
            if totals is None:
                totals = [value.copy() for value in values]
            else:
                for total, value in zip(totals, values):
                    total += value
        s0, s1, s2, event_count, event_x_sum, event_eta_sum = totals
        if np.any(s0 <= 0):
            raise FloatingPointError("non-positive distributed Cox risk-set denominator")

        means = s1 / s0[:, None]
        score = np.sum(event_x_sum - event_count[:, None] * means, axis=0)
        information = np.zeros((len(beta), len(beta)), dtype=np.float64)
        for index in range(len(event_times)):
            covariance = s2[index] / s0[index] - np.outer(means[index], means[index])
            information += event_count[index] * covariance
        log_likelihood = float(
            np.sum(
                event_eta_sum
                - event_count * (global_maxima + np.log(s0))
            )
        )

        l2 = float(self.config.webdisco_l2) if self.config is not None else 0.0
        if l2:
            log_likelihood -= 0.5 * l2 * float(beta @ beta)
            score -= l2 * beta
            information += l2 * np.eye(len(beta))

        n_clients = len(clients)
        n_times = len(event_times)
        n_features = len(beta)
        # Analytical float64 payload count. It excludes serialization and
        # protocol framing and is therefore labelled as a payload estimate.
        first_pass = 2 * n_clients * n_times
        second_pass = n_clients * n_times * (
            2 + 2 * n_features + n_features * n_features
        )
        coefficient_broadcast = n_clients * n_features
        payload_bytes = int(8 * (first_pass + second_pass + coefficient_broadcast))
        return log_likelihood, score, information, payload_bytes

    def fit_coefficients(self, client_sets) -> WebDISCOFitResult:
        clients = self._coerce_clients(client_sets)
        n_features = next(iter(clients.values()))[0].shape[1]
        event_times = self.event_times(clients)
        beta = np.zeros(n_features, dtype=np.float64)
        beta_history = []
        likelihood_history = []
        total_bytes = 0
        converged = False
        final_score_norm = float("inf")

        for _ in range(int(self.config.webdisco_max_iter)):
            likelihood, score, information, payload = self.distributed_objective(
                clients, beta, event_times
            )
            total_bytes += payload
            final_score_norm = float(np.linalg.norm(score))
            ridge = float(self.config.webdisco_ridge)
            try:
                step = np.linalg.solve(
                    information + ridge * np.eye(n_features), score
                )
            except np.linalg.LinAlgError:
                step = np.linalg.pinv(
                    information + ridge * np.eye(n_features)
                ) @ score

            step_norm = float(np.linalg.norm(step))
            max_step = float(self.config.webdisco_max_step_norm)
            if step_norm > max_step:
                step *= max_step / step_norm

            scale = 1.0
            accepted = False
            candidate_likelihood = likelihood
            while scale >= 2.0 ** -12:
                candidate = beta + scale * step
                candidate_likelihood, _, _, candidate_payload = self.distributed_objective(
                    clients, candidate, event_times
                )
                total_bytes += candidate_payload
                if candidate_likelihood >= likelihood - 1e-12:
                    beta = candidate
                    accepted = True
                    break
                scale *= 0.5
            if not accepted:
                break

            beta_history.append(beta.copy())
            likelihood_history.append(float(candidate_likelihood))
            if scale * step_norm <= float(self.config.webdisco_tolerance):
                converged = True
                break

        final_likelihood, final_score, _, payload = self.distributed_objective(
            clients, beta, event_times
        )
        total_bytes += payload
        final_score_norm = float(np.linalg.norm(final_score))
        if final_score_norm <= max(float(self.config.webdisco_tolerance), 1e-6):
            converged = True
        if not beta_history:
            beta_history.append(beta.copy())
            likelihood_history.append(float(final_likelihood))
        return WebDISCOFitResult(
            coefficients=beta.copy(),
            coefficient_history=tuple(beta_history),
            log_likelihood_history=tuple(likelihood_history),
            score_norm=final_score_norm,
            converged=converged,
            n_iterations=len(beta_history),
            communication_bytes=total_bytes,
        )

    def aggregate(self, *args, **kwargs):
        raise RuntimeError("WebDISCO-style owns its risk-set-statistics training loop")

    def metadata(self):
        result = super().metadata()
        result.update(
            {
                "implementation_scope": (
                    "independent Breslow linear-Cox risk-set-statistics implementation; "
                    "not the original WebDISCO web service"
                ),
                "max_iterations": int(self.config.webdisco_max_iter),
                "tolerance": float(self.config.webdisco_tolerance),
                "l2": float(self.config.webdisco_l2),
            }
        )
        return result

distributed_objective(client_sets, beta, event_times=None)

Return pooled Breslow log likelihood, score, information and bytes.

Source code in federated_survival/protocols/webdisco.py

def distributed_objective(self, client_sets, beta, event_times=None):
    """Return pooled Breslow log likelihood, score, information and bytes."""
    clients = self._coerce_clients(client_sets)
    beta = np.asarray(beta, dtype=np.float64).reshape(-1)
    if event_times is None:
        event_times = self.event_times(clients)
    event_times = np.asarray(event_times, dtype=np.float64)
    local_maxima = np.stack(
        [self._risk_maxima(x, y, beta, event_times) for x, y in clients.values()]
    )
    global_maxima = np.max(local_maxima, axis=0)
    if not np.isfinite(global_maxima).all():
        raise FloatingPointError("an event time has an empty global risk set")

    totals = None
    for x, y in clients.values():
        values = self._local_statistics(x, y, beta, event_times, global_maxima)
        if totals is None:
            totals = [value.copy() for value in values]
        else:
            for total, value in zip(totals, values):
                total += value
    s0, s1, s2, event_count, event_x_sum, event_eta_sum = totals
    if np.any(s0 <= 0):
        raise FloatingPointError("non-positive distributed Cox risk-set denominator")

    means = s1 / s0[:, None]
    score = np.sum(event_x_sum - event_count[:, None] * means, axis=0)
    information = np.zeros((len(beta), len(beta)), dtype=np.float64)
    for index in range(len(event_times)):
        covariance = s2[index] / s0[index] - np.outer(means[index], means[index])
        information += event_count[index] * covariance
    log_likelihood = float(
        np.sum(
            event_eta_sum
            - event_count * (global_maxima + np.log(s0))
        )
    )

    l2 = float(self.config.webdisco_l2) if self.config is not None else 0.0
    if l2:
        log_likelihood -= 0.5 * l2 * float(beta @ beta)
        score -= l2 * beta
        information += l2 * np.eye(len(beta))

    n_clients = len(clients)
    n_times = len(event_times)
    n_features = len(beta)
    # Analytical float64 payload count. It excludes serialization and
    # protocol framing and is therefore labelled as a payload estimate.
    first_pass = 2 * n_clients * n_times
    second_pass = n_clients * n_times * (
        2 + 2 * n_features + n_features * n_features
    )
    coefficient_broadcast = n_clients * n_features
    payload_bytes = int(8 * (first_pass + second_pass + coefficient_broadcast))
    return log_likelihood, score, information, payload_bytes

Backward-compatible core

Configuration class for federated survival analysis.

Source code in federated_survival/core/config.py

@dataclass
class FSAConfig:
    """Configuration class for federated survival analysis."""

    # Data parameters
    dataset_name: Optional[str] = None  # Dataset name when using real data
    mode: str = "simulate"  # Options: 'real', 'simulate'

    # Data configuration
    n_samples: int = 1000  # Number of samples
    n_features: int = 20  # Number of features
    censor_rate: float = 0.4  # Censoring rate
    data_type: str = "weibull"  # Options: 'weibull', 'lognormal', '1', '2', '3', '4'
    sim_type: str = "1"  # Simulation data type

    # Model parameters
    model_type: str = "PC-Hazard"  # Model type
    num_nodes: Tuple[int, ...] = (32, 32)  # Number of layers and nodes per layer
    num_durations: int = 25  # Number of discretized time intervals
    batch_norm: bool = False  # Whether to use batch normalization
    dropout: float = 0.1  # Dropout rate
    activation: nn.Module = nn.ReLU  # Activation function

    # Federated learning parameters
    num_clients: int = 5  # Number of clients
    global_epochs: int = 50  # Number of global training rounds
    early_stopping: bool = False  # Whether to use early stopping
    early_stopping_patience: int = 10  # Early stopping patience (rounds)
    local_epochs: int = 1  # Local update steps E before each aggregation (legacy field name)
    batch_size: int = 32  # Batch size
    full_batch: bool = False  # False for independent mini-batch local SGD; True only as a full-batch special case
    learning_rate: float = 1e-3  # Learning rate
    optimizer: str = "adam"  # Options: 'adam', 'sgd'
    weight_decay: float = 0.05  # Optimizer weight decay
    client_sample_ratio: float = 1.0  # Fraction of clients selected per round
    split_method: str = "iid"  # Data split method, options: 'iid', 'non-iid', 'time-non-iid'

    # Federated protocol parameters
    federated_protocol: str = "FedAvg"  # FedAvg, FedProx, FedOpt, WebDISCO-style
    proximal_mu: float = 0.01  # FedProx proximal term coefficient
    server_optimizer: str = "adam"  # FedOpt: adam, yogi, adagrad
    server_learning_rate: float = 0.01
    server_beta1: float = 0.9
    server_beta2: float = 0.99
    server_tau: float = 1e-3
    webdisco_max_iter: int = 50
    webdisco_tolerance: float = 1e-7
    webdisco_l2: float = 0.0
    webdisco_ridge: float = 1e-8
    webdisco_max_step_norm: float = 5.0

    # Data augmentation parameters
    k: float = 0.5  # Proportion of augmented data
    latent_num: int = 10
    hidden_num: int = 30
    alpha: float = 1.0
    beta: float = 1.0
    augmentation_sampling: str = "sparse"
    augmentation_sparse_gamma: float = 0.1

    # Differential privacy parameters
    use_differential_privacy: bool = False  # Whether to use differential privacy
    dp_mechanism: str = "gaussian"  # Differential privacy mechanism: 'gaussian', 'laplace', 'exponential'
    dp_epsilon: float = 1.0  # Privacy budget (epsilon)
    dp_delta: float = 1e-5  # Failure probability (delta) - required only for the Gaussian mechanism
    dp_sensitivity: float = 1.0  # Sensitivity
    dp_noise_multiplier: float = 1.0  # Noise multiplier - used only by the Gaussian mechanism
    dp_clip_norm: float = 1.0  # Gradient clipping norm

    # Other parameters
    verbose: bool = False  # Whether to print verbose output
    show_progress: bool = False  # Whether to show the tqdm progress bar
    random_seed: int = 42  # Random seed
    evaluation_quantiles: Tuple[float, float] = (0.05, 0.95)
    cox_loss_normalization: str = "patient"  # 'event' reproduces raw pycox scaling
    prediction_validation: str = "raise"  # 'raise', 'warn', or 'none'

    # Base configuration
    n_rounds: int = 100
    random_state: Optional[int] = None

    # Model configuration
    model_params: Dict[str, Any] = field(default_factory=dict)

    # Data split configuration
    split_alpha: float = 0.5  # Dirichlet distribution parameter for non-iid splits
    test_size: float = 0.2

    def __post_init__(self):
        """Validate configuration parameters."""
        # Validate mode
        valid_modes = ["real", "simulate"]
        if self.mode not in valid_modes:
            raise ValueError(f"mode must be one of {valid_modes}")

        # A dataset name is required in 'real' mode
        if self.mode == "real" and self.dataset_name is None:
            raise ValueError("dataset_name must be provided when mode is 'real'")

        if not 0 <= self.censor_rate <= 1:
            raise ValueError("censor_rate must be between 0 and 1")

        if not 0 < self.client_sample_ratio <= 1:
            raise ValueError("client_sample_ratio must be between 0 and 1")

        self.optimizer = self.optimizer.lower()
        if self.optimizer not in ["adam", "sgd"]:
            raise ValueError("optimizer must be 'adam' or 'sgd'")
        if self.weight_decay < 0:
            raise ValueError("weight_decay must be non-negative")
        if self.num_clients <= 0:
            raise ValueError("num_clients must be positive")
        if self.global_epochs <= 0 or self.local_epochs <= 0:
            raise ValueError("global_epochs and local_epochs must be positive")
        if self.batch_size <= 0:
            raise ValueError("batch_size must be positive")
        if self.learning_rate <= 0:
            raise ValueError("learning_rate must be positive")

        from ..protocols import get_federated_protocol

        protocol = get_federated_protocol(self.federated_protocol)
        self.federated_protocol = protocol.name
        if self.proximal_mu < 0:
            raise ValueError("proximal_mu must be non-negative")
        self.server_optimizer = self.server_optimizer.lower()
        if self.server_optimizer not in ["adam", "yogi", "adagrad"]:
            raise ValueError("server_optimizer must be 'adam', 'yogi', or 'adagrad'")
        if self.server_learning_rate <= 0:
            raise ValueError("server_learning_rate must be positive")
        if not 0 <= self.server_beta1 < 1 or not 0 <= self.server_beta2 < 1:
            raise ValueError("server_beta1 and server_beta2 must be in [0, 1)")
        if self.server_tau <= 0:
            raise ValueError("server_tau must be positive")
        if self.webdisco_max_iter <= 0:
            raise ValueError("webdisco_max_iter must be positive")
        if self.webdisco_tolerance <= 0:
            raise ValueError("webdisco_tolerance must be positive")
        if self.webdisco_l2 < 0 or self.webdisco_ridge < 0:
            raise ValueError("webdisco_l2 and webdisco_ridge must be non-negative")
        if self.webdisco_max_step_norm <= 0:
            raise ValueError("webdisco_max_step_norm must be positive")
        self.cox_loss_normalization = self.cox_loss_normalization.lower()
        if self.cox_loss_normalization not in ["patient", "event"]:
            raise ValueError("cox_loss_normalization must be 'patient' or 'event'")
        self.prediction_validation = self.prediction_validation.lower()
        if self.prediction_validation not in ["raise", "warn", "none"]:
            raise ValueError("prediction_validation must be 'raise', 'warn', or 'none'")
        q_low, q_high = self.evaluation_quantiles
        if not 0 <= q_low < q_high <= 1:
            raise ValueError("evaluation_quantiles must satisfy 0 <= low < high <= 1")

        # Validate model type
        # Resolve model support through the public registry so third-party
        # adapters can be configured without editing this dataclass.
        from ..models import available_model_adapters

        valid_model_types = list(available_model_adapters())
        if self.model_type not in valid_model_types:
            raise ValueError(f"model_type must be one of {valid_model_types}")

        # Validate intermediate layer parameters for CoxPH
        if self.model_type == "CoxPH":
            if self.num_nodes != ():
                raise ValueError("num_nodes must be () for CoxPH")

        protocol.validate_config(self)

        # Validate data split method
        self.split_method = self.split_method.lower()
        valid_split_methods = [
            "iid",
            "random",
            "non-iid",
            "censoring-non-iid",
            "time-non-iid",
            "dirichlet",
        ]
        if self.split_method not in valid_split_methods:
            raise ValueError(f"split_method must be one of {valid_split_methods}")

        # Validate numerical ranges
        if not 0 < self.test_size < 1:
            raise ValueError("test_size must be between 0 and 1")
        if not 0 < self.k < 1:
            raise ValueError("k must be between 0 and 1")

        # Validate data augmentation parameters
        if self.latent_num <= 0:
            raise ValueError("latent_num must be positive")
        if self.hidden_num <= 0:
            raise ValueError("hidden_num must be positive")
        self.augmentation_sampling = self.augmentation_sampling.lower().replace("-", "_")
        if self.augmentation_sampling not in ["sparse", "unconditional"]:
            raise ValueError("augmentation_sampling must be 'sparse' or 'unconditional'")
        if self.augmentation_sparse_gamma < 0:
            raise ValueError("augmentation_sparse_gamma must be non-negative")
        if self.alpha <= 0:
            raise ValueError("alpha must be positive")
        if self.beta <= 0:
            raise ValueError("beta must be positive")

        # Validate num_durations
        if self.num_durations <= 0:
            raise ValueError("num_durations must be positive")

        # Validate dropout
        if not 0 <= self.dropout <= 1:
            raise ValueError("dropout must be between 0 and 1")

        # Validate n_rounds
        if self.n_rounds <= 0:
            raise ValueError("n_rounds must be positive")

        # Validate split_alpha
        if self.split_alpha <= 0:
            raise ValueError("split_alpha must be positive")

        # Validate differential privacy parameters
        if self.use_differential_privacy:
            # Validate mechanism type
            valid_dp_mechanisms = ["gaussian", "laplace", "exponential"]
            if self.dp_mechanism not in valid_dp_mechanisms:
                raise ValueError(f"dp_mechanism must be one of {valid_dp_mechanisms}")

            # Validate common parameters
            if self.dp_epsilon <= 0:
                raise ValueError("dp_epsilon must be positive")
            if self.dp_sensitivity <= 0:
                raise ValueError("dp_sensitivity must be positive")

            # Gaussian mechanism specific parameters
            if self.dp_mechanism == "gaussian":
                if not 0 < self.dp_delta < 1:
                    raise ValueError("dp_delta must be between 0 and 1 for Gaussian mechanism")
                if self.dp_noise_multiplier <= 0:
                    raise ValueError("dp_noise_multiplier must be positive for Gaussian mechanism")

            # Validate gradient clipping parameters (applies to Gaussian and Laplace mechanisms)
            if self.dp_mechanism in ["gaussian", "laplace"]:
                if self.dp_clip_norm <= 0:
                    raise ValueError("dp_clip_norm must be positive")

            # Warn when a configured DP field cannot influence the selected
            # mechanism.  The Gaussian noise scale is driven by
            # dp_noise_multiplier and the Laplace scale by dp_epsilon, so
            # switching dp_mechanism turns one of them into an inert value that
            # looks like a working privacy control.  dp_delta never affects the
            # noise scale either.  See core.differential_privacy for the
            # authoritative mapping.
            from .differential_privacy import DP_MECHANISM_INERT_KNOBS

            inert = DP_MECHANISM_INERT_KNOBS[self.dp_mechanism]
            driver = (
                "dp_noise_multiplier" if self.dp_mechanism == "gaussian" else "dp_epsilon"
            )
            self._warn_inactive_dp_parameters(inert, driver)

        # Set default model parameters
        if not self.model_params:
            if self.model_type == "PC-Hazard":
                self.model_params = {"n_intervals": 10, "hidden_size": 32, "dropout": 0.1}
            elif self.model_type == "LogisticHazard":
                self.model_params = {"n_intervals": 10, "hidden_size": 32, "dropout": 0.1}
            elif self.model_type == "DeepHit":
                self.model_params = {"n_intervals": 10, "hidden_size": 32, "dropout": 0.1}
            elif self.model_type == "CoxTime":
                self.model_params = {"hidden_size": 32, "dropout": 0.1}
            elif self.model_type in ["DeepSurv", "CoxPH"]:
                self.model_params = {"l2_reg": 0.01}
            elif self.model_type == "CoxCC":
                self.model_params = {"l2_reg": 0.01}

    def _warn_inactive_dp_parameters(self, inert, driver) -> None:
        """Emit one warning per inert DP field that was set away from its default.

        A field left at its default cannot mislead anyone, so it is silent.  The
        warning fires only when the user has actually chosen a non-default value
        for a parameter that the selected mechanism ignores.
        """
        import warnings
        from dataclasses import MISSING, fields as dataclass_fields

        defaults = {}
        for f in dataclass_fields(self):
            if f.default is not MISSING:
                defaults[f.name] = f.default
            elif f.default_factory is not MISSING:  # type: ignore[misc]
                defaults[f.name] = f.default_factory()  # type: ignore[misc]

        for name in inert:
            value = getattr(self, name, None)
            default = defaults.get(name)
            if default is not None and value == default:
                # Untouched default: not an active misconfiguration.
                continue
            warnings.warn(
                f"dp_mechanism='{self.dp_mechanism}' does not use {name} "
                f"(currently {value!r}); the injected noise is controlled by "
                f"{driver}. Setting {name} has no effect on the noise added to "
                f"model updates. Switch dp_mechanism or adjust {driver} instead.",
                UserWarning,
                stacklevel=4,
            )

__post_init__()

Validate configuration parameters.

Source code in federated_survival/core/config.py

def __post_init__(self):
    """Validate configuration parameters."""
    # Validate mode
    valid_modes = ["real", "simulate"]
    if self.mode not in valid_modes:
        raise ValueError(f"mode must be one of {valid_modes}")

    # A dataset name is required in 'real' mode
    if self.mode == "real" and self.dataset_name is None:
        raise ValueError("dataset_name must be provided when mode is 'real'")

    if not 0 <= self.censor_rate <= 1:
        raise ValueError("censor_rate must be between 0 and 1")

    if not 0 < self.client_sample_ratio <= 1:
        raise ValueError("client_sample_ratio must be between 0 and 1")

    self.optimizer = self.optimizer.lower()
    if self.optimizer not in ["adam", "sgd"]:
        raise ValueError("optimizer must be 'adam' or 'sgd'")
    if self.weight_decay < 0:
        raise ValueError("weight_decay must be non-negative")
    if self.num_clients <= 0:
        raise ValueError("num_clients must be positive")
    if self.global_epochs <= 0 or self.local_epochs <= 0:
        raise ValueError("global_epochs and local_epochs must be positive")
    if self.batch_size <= 0:
        raise ValueError("batch_size must be positive")
    if self.learning_rate <= 0:
        raise ValueError("learning_rate must be positive")

    from ..protocols import get_federated_protocol

    protocol = get_federated_protocol(self.federated_protocol)
    self.federated_protocol = protocol.name
    if self.proximal_mu < 0:
        raise ValueError("proximal_mu must be non-negative")
    self.server_optimizer = self.server_optimizer.lower()
    if self.server_optimizer not in ["adam", "yogi", "adagrad"]:
        raise ValueError("server_optimizer must be 'adam', 'yogi', or 'adagrad'")
    if self.server_learning_rate <= 0:
        raise ValueError("server_learning_rate must be positive")
    if not 0 <= self.server_beta1 < 1 or not 0 <= self.server_beta2 < 1:
        raise ValueError("server_beta1 and server_beta2 must be in [0, 1)")
    if self.server_tau <= 0:
        raise ValueError("server_tau must be positive")
    if self.webdisco_max_iter <= 0:
        raise ValueError("webdisco_max_iter must be positive")
    if self.webdisco_tolerance <= 0:
        raise ValueError("webdisco_tolerance must be positive")
    if self.webdisco_l2 < 0 or self.webdisco_ridge < 0:
        raise ValueError("webdisco_l2 and webdisco_ridge must be non-negative")
    if self.webdisco_max_step_norm <= 0:
        raise ValueError("webdisco_max_step_norm must be positive")
    self.cox_loss_normalization = self.cox_loss_normalization.lower()
    if self.cox_loss_normalization not in ["patient", "event"]:
        raise ValueError("cox_loss_normalization must be 'patient' or 'event'")
    self.prediction_validation = self.prediction_validation.lower()
    if self.prediction_validation not in ["raise", "warn", "none"]:
        raise ValueError("prediction_validation must be 'raise', 'warn', or 'none'")
    q_low, q_high = self.evaluation_quantiles
    if not 0 <= q_low < q_high <= 1:
        raise ValueError("evaluation_quantiles must satisfy 0 <= low < high <= 1")

    # Validate model type
    # Resolve model support through the public registry so third-party
    # adapters can be configured without editing this dataclass.
    from ..models import available_model_adapters

    valid_model_types = list(available_model_adapters())
    if self.model_type not in valid_model_types:
        raise ValueError(f"model_type must be one of {valid_model_types}")

    # Validate intermediate layer parameters for CoxPH
    if self.model_type == "CoxPH":
        if self.num_nodes != ():
            raise ValueError("num_nodes must be () for CoxPH")

    protocol.validate_config(self)

    # Validate data split method
    self.split_method = self.split_method.lower()
    valid_split_methods = [
        "iid",
        "random",
        "non-iid",
        "censoring-non-iid",
        "time-non-iid",
        "dirichlet",
    ]
    if self.split_method not in valid_split_methods:
        raise ValueError(f"split_method must be one of {valid_split_methods}")

    # Validate numerical ranges
    if not 0 < self.test_size < 1:
        raise ValueError("test_size must be between 0 and 1")
    if not 0 < self.k < 1:
        raise ValueError("k must be between 0 and 1")

    # Validate data augmentation parameters
    if self.latent_num <= 0:
        raise ValueError("latent_num must be positive")
    if self.hidden_num <= 0:
        raise ValueError("hidden_num must be positive")
    self.augmentation_sampling = self.augmentation_sampling.lower().replace("-", "_")
    if self.augmentation_sampling not in ["sparse", "unconditional"]:
        raise ValueError("augmentation_sampling must be 'sparse' or 'unconditional'")
    if self.augmentation_sparse_gamma < 0:
        raise ValueError("augmentation_sparse_gamma must be non-negative")
    if self.alpha <= 0:
        raise ValueError("alpha must be positive")
    if self.beta <= 0:
        raise ValueError("beta must be positive")

    # Validate num_durations
    if self.num_durations <= 0:
        raise ValueError("num_durations must be positive")

    # Validate dropout
    if not 0 <= self.dropout <= 1:
        raise ValueError("dropout must be between 0 and 1")

    # Validate n_rounds
    if self.n_rounds <= 0:
        raise ValueError("n_rounds must be positive")

    # Validate split_alpha
    if self.split_alpha <= 0:
        raise ValueError("split_alpha must be positive")

    # Validate differential privacy parameters
    if self.use_differential_privacy:
        # Validate mechanism type
        valid_dp_mechanisms = ["gaussian", "laplace", "exponential"]
        if self.dp_mechanism not in valid_dp_mechanisms:
            raise ValueError(f"dp_mechanism must be one of {valid_dp_mechanisms}")

        # Validate common parameters
        if self.dp_epsilon <= 0:
            raise ValueError("dp_epsilon must be positive")
        if self.dp_sensitivity <= 0:
            raise ValueError("dp_sensitivity must be positive")

        # Gaussian mechanism specific parameters
        if self.dp_mechanism == "gaussian":
            if not 0 < self.dp_delta < 1:
                raise ValueError("dp_delta must be between 0 and 1 for Gaussian mechanism")
            if self.dp_noise_multiplier <= 0:
                raise ValueError("dp_noise_multiplier must be positive for Gaussian mechanism")

        # Validate gradient clipping parameters (applies to Gaussian and Laplace mechanisms)
        if self.dp_mechanism in ["gaussian", "laplace"]:
            if self.dp_clip_norm <= 0:
                raise ValueError("dp_clip_norm must be positive")

        # Warn when a configured DP field cannot influence the selected
        # mechanism.  The Gaussian noise scale is driven by
        # dp_noise_multiplier and the Laplace scale by dp_epsilon, so
        # switching dp_mechanism turns one of them into an inert value that
        # looks like a working privacy control.  dp_delta never affects the
        # noise scale either.  See core.differential_privacy for the
        # authoritative mapping.
        from .differential_privacy import DP_MECHANISM_INERT_KNOBS

        inert = DP_MECHANISM_INERT_KNOBS[self.dp_mechanism]
        driver = (
            "dp_noise_multiplier" if self.dp_mechanism == "gaussian" else "dp_epsilon"
        )
        self._warn_inactive_dp_parameters(inert, driver)

    # Set default model parameters
    if not self.model_params:
        if self.model_type == "PC-Hazard":
            self.model_params = {"n_intervals": 10, "hidden_size": 32, "dropout": 0.1}
        elif self.model_type == "LogisticHazard":
            self.model_params = {"n_intervals": 10, "hidden_size": 32, "dropout": 0.1}
        elif self.model_type == "DeepHit":
            self.model_params = {"n_intervals": 10, "hidden_size": 32, "dropout": 0.1}
        elif self.model_type == "CoxTime":
            self.model_params = {"hidden_size": 32, "dropout": 0.1}
        elif self.model_type in ["DeepSurv", "CoxPH"]:
            self.model_params = {"l2_reg": 0.01}
        elif self.model_type == "CoxCC":
            self.model_params = {"l2_reg": 0.01}

Runner for federated survival analysis.

Source code in federated_survival/core/runner.py

class FSARunner:
    """Runner for federated survival analysis."""

    def __init__(self, config: FSAConfig):
        """
        Initialize the federated survival analysis runner.

        Args:
            config: Federated learning configuration.
        """
        self.config = config
        self.adapter = get_model_adapter(config.model_type)
        self.protocol = get_federated_protocol(config.federated_protocol)
        self.protocol.validate_config(config)
        self.set_random_seed()

    def set_random_seed(self):
        """Set the random seed."""
        if self.config.random_seed is not None:
            random.seed(self.config.random_seed)
            np.random.seed(self.config.random_seed)
            torch.manual_seed(self.config.random_seed)
            if torch.cuda.is_available():
                torch.cuda.manual_seed(self.config.random_seed)

    def _get_label_transform(self):
        """Get the label transformer."""
        self.adapter = get_model_adapter(self.config.model_type)
        return self.adapter.create_label_transform(self.config)

    def _get_model(self, net: nn.Module, labtrans=None):
        """Get the survival analysis model."""
        self.adapter = get_model_adapter(self.config.model_type)
        return self.adapter.build_model(
            net,
            self.config,
            label_transform=labtrans,
            optimizer=torch.optim.Adam,
        )

    def run(
        self, data: Any, type: str = "raw", aug_method: str = "MVAEC"
    ) -> Dict[str, List[float]]:
        """
        Run federated survival analysis.

        Args:
            data: Survival data with the following attributes:
                - clients_set: Client datasets in the form {client_id: (train_X, train_y)}
                - raw_aug_clients_set: Client datasets augmented from raw data,
                  in the form {client_id: (train_X, train_y)}
                - test_data: Test data
                - test_label: Test labels
            type: Data type, 'raw' or 'raw_aug'
            aug_method: Data augmentation method, 'MVAEC' or 'MVAES'

        Returns:
            dict: {'train_Cindex': train_Cindex, 'train_IBS': train_IBS, 'test_Cindex': test_Cindex, 'test_IBS': test_IBS}
        """
        self.adapter = get_model_adapter(self.config.model_type)

        # Get training data
        if type == "raw":
            clients_set = data.clients_set
        elif type == "raw_aug":
            augmenter = DataAugmenter(
                latent_num=self.config.latent_num,
                hidden_num=self.config.hidden_num,
                alpha=self.config.alpha,
                beta=self.config.beta,
                sampling=self.config.augmentation_sampling,
                sparse_gamma=self.config.augmentation_sparse_gamma,
            )
            if aug_method == "MVAEC":
                clients_set = augmenter.mvaec(data.clients_set, k=self.config.k)
            elif aug_method == "MVAES":
                clients_set = augmenter.mvaes(data.clients_set, k=self.config.k)
            else:
                raise ValueError(f"Unsupported augmentation method: {aug_method}")
        else:
            raise ValueError(f"Unsupported data type: {type}")

        if self.protocol.owns_training_loop:
            if type != "raw":
                raise ValueError(
                    f"protocol {self.protocol.name!r} does not support data augmentation"
                )
            return self._run_webdisco_style(data, clients_set)

        _y_train = np.empty(shape=(0, 2), dtype=np.float32)
        for i in clients_set.keys():
            _y_train = np.vstack((_y_train, clients_set[i][1]))
        y_train = self._get_target(_y_train)

        # Let the registered model adapter own label fitting and output shape.
        labtrans = self.adapter.configure_targets(self.config, *y_train)

        # Get test data
        durations_test, events_test = self._get_target(data.test_label)

        # Initialize server and clients
        server = Server(self.config)
        self.protocol.begin_run(self.config, server.global_model.state_dict())
        clients = []
        for i in clients_set.keys():
            clients.append(
                Client(
                    self.config,
                    server.global_model,
                    clients_set[i],
                    i,
                    protocol=self.protocol,
                )
            )

        # If differential privacy is enabled, print privacy protection info
        if self.config.use_differential_privacy:
            if self.config.verbose:
                mechanism = (
                    self.config.dp_mechanism if hasattr(self.config, "dp_mechanism") else "gaussian"
                )
                print("Differential privacy enabled:")
                print(f"  - mechanism: {mechanism.upper()}")
                print(f"  - privacy budget (epsilon): {self.config.dp_epsilon}")
                if mechanism == "gaussian":
                    print(f"  - failure probability (delta): {self.config.dp_delta}")
                    print(f"  - noise multiplier: {self.config.dp_noise_multiplier}")
                print(f"  - sensitivity: {self.config.dp_sensitivity}")
                if mechanism in ["gaussian", "laplace"]:
                    print(f"  - gradient clipping norm: {self.config.dp_clip_norm}")

        # Record metrics
        train_Cindex = []
        train_IBS = []
        test_Cindex = []
        test_IBS = []

        # Federated learning training
        train_loss = []
        update_direction_norm = []
        client_drift = []
        update_direction_dispersion = []
        selected_clients = []
        communication_bytes = []
        model_bytes = sum(
            value.numel() * value.element_size()
            for value in server.global_model.state_dict().values()
        )

        for e in tqdm(
            range(self.config.global_epochs),
            disable=not self.config.show_progress,
        ):
            # Select clients
            candidates = random.sample(
                clients, max(round(self.config.client_sample_ratio * len(clients)), 1)
            )

            # Compute the total number of samples
            N = sum(c.N for c in candidates)
            round_directions = []
            round_weights = []
            round_losses = []
            round_drifts = []
            round_states = []
            global_state = {
                name: value.detach().clone()
                for name, value in server.global_model.state_dict().items()
            }

            # Client training
            for i, c in enumerate(candidates, 1):
                if self.config.verbose:
                    print(f"Epoch {e+1}/{self.config.global_epochs}, Client {i}/{len(candidates)}")

                local = c.local_train(server.global_model, e)
                round_directions.append(c.last_update_direction)
                round_weights.append(c.N / N)
                round_losses.append(c.last_loss)
                round_drifts.append(c.last_drift_norm)
                round_states.append(
                    {name: value.detach().clone() for name, value in local.state_dict().items()}
                )

            # Server update
            aggregated_state = self.protocol.aggregate(
                global_state,
                round_states,
                round_weights,
                round_index=e,
            )
            server.model_update(aggregated_state, len(candidates))

            weights = torch.tensor(round_weights, dtype=torch.float32)
            directions = torch.stack(round_directions)
            mean_direction = torch.sum(weights[:, None] * directions, dim=0)
            heterogeneity = torch.sum(
                weights * torch.sum((directions - mean_direction[None, :]) ** 2, dim=1)
            )
            train_loss.append(float(np.average(round_losses, weights=round_weights)))
            update_direction_norm.append(float(torch.linalg.vector_norm(mean_direction)))
            client_drift.append(float(np.average(round_drifts, weights=round_weights)))
            update_direction_dispersion.append(float(heterogeneity))
            selected_clients.append([str(c.client_id) for c in candidates])
            communication_bytes.append(
                self.protocol.round_communication_bytes(len(candidates), model_bytes)
            )

            # Evaluate the model
            index = server.model_eval(clients_set, labtrans)
            train_Cindex.append(index[0])
            train_IBS.append(index[1])

            # Compute test metrics
            model = self._get_model(server.global_model, labtrans)

            pooled_x = np.concatenate(
                [client_data[0] for client_data in clients_set.values()], axis=0
            )
            pooled_y = np.concatenate(
                [client_data[1] for client_data in clients_set.values()], axis=0
            )
            self.adapter.prepare_for_prediction(model, pooled_x, pooled_y, labtrans)
            surv, prediction_diagnostics = self.adapter.predict_survival(
                model,
                data.test_data,
                validation_mode=self.config.prediction_validation,
            )
            ev = EvalSurv(surv, durations_test, events_test, censor_surv="km")
            time_grid = evaluation_time_grid(
                durations_test,
                surv.index.values,
                self.config.evaluation_quantiles,
            )
            test_Cindex.append(ev.concordance_td())
            test_IBS.append(ev.integrated_brier_score(time_grid))

            if self.config.early_stopping:
                # Early stopping
                patience = self.config.early_stopping_patience
                if (
                    len(train_Cindex) > patience
                    and max(train_Cindex[-patience:]) <= train_Cindex[-patience - 1]
                ):
                    break

        self.server_ = server
        self.model_ = model
        self.labtrans_ = labtrans
        self.training_features_ = pooled_x
        self.training_labels_ = pooled_y
        self.last_prediction_diagnostics_ = prediction_diagnostics

        return {
            "train_Cindex": train_Cindex,
            "train_IBS": train_IBS,
            "test_Cindex": test_Cindex,
            "test_IBS": test_IBS,
            "train_loss": train_loss,
            "update_direction_norm": update_direction_norm,
            "client_drift": client_drift,
            "update_direction_dispersion": update_direction_dispersion,
            # Backward-compatible alias. This is not the common-point B^2 in
            # the convergence assumption when local_epochs > 1.
            "empirical_heterogeneity": update_direction_dispersion,
            "selected_clients": selected_clients,
            "communication_bytes": communication_bytes,
            "federated_protocol": self.protocol.name,
            "protocol_metadata": self.protocol.metadata(),
            "telemetry_is_path_average_gradient": (
                self.config.optimizer == "sgd"
                and self.config.weight_decay == 0
                and not self.config.use_differential_privacy
            ),
            "telemetry_is_average_gradient": (
                self.config.optimizer == "sgd"
                and self.config.weight_decay == 0
                and not self.config.use_differential_privacy
            ),
            # With one local SGD step, every recorded direction is evaluated at
            # the broadcast point.  It equals the deterministic local-objective
            # gradient only in full-batch mode; otherwise it is a stochastic
            # gradient at that common point.
            "telemetry_is_common_point_gradient": (
                self.config.optimizer == "sgd"
                and self.config.weight_decay == 0
                and not self.config.use_differential_privacy
                and self.config.local_epochs == 1
            ),
            "telemetry_is_common_point_full_gradient": (
                self.config.optimizer == "sgd"
                and self.config.weight_decay == 0
                and not self.config.use_differential_privacy
                and self.config.local_epochs == 1
                and self.config.full_batch
            ),
        }

    def _run_webdisco_style(self, data, clients_set):
        """Run the protocol-owned distributed linear Cox Newton solver."""
        self.adapter = get_model_adapter(self.config.model_type)
        pooled_x = np.concatenate(
            [client_data[0] for client_data in clients_set.values()], axis=0
        ).astype(np.float32)
        pooled_y = np.concatenate(
            [client_data[1] for client_data in clients_set.values()], axis=0
        ).astype(np.float32)
        durations_train, events_train = self._get_target(pooled_y)
        durations_test, events_test = self._get_target(data.test_label)
        labtrans = self.adapter.configure_targets(self.config, durations_train, events_train)

        server = Server(self.config)
        self.protocol.begin_run(self.config, server.global_model.state_dict())
        fit_result = self.protocol.fit_coefficients(clients_set)

        state = {
            name: value.detach().clone() for name, value in server.global_model.state_dict().items()
        }
        parameter_items = [
            (name, value)
            for name, value in state.items()
            if value.is_floating_point() and value.numel() == self.config.n_features
        ]
        if len(parameter_items) != 1:
            raise RuntimeError(
                "WebDISCO-style requires one bias-free linear CoxPH parameter tensor"
            )
        parameter_name, parameter_value = parameter_items[0]
        state[parameter_name] = torch.as_tensor(
            fit_result.coefficients,
            dtype=parameter_value.dtype,
            device=parameter_value.device,
        ).reshape(parameter_value.shape)
        server.model_update(state)

        model = self._get_model(server.global_model, labtrans)
        self.adapter.prepare_for_prediction(model, pooled_x, pooled_y, labtrans)
        train_survival, train_diagnostics = self.adapter.predict_survival(
            model,
            pooled_x,
            validation_mode=self.config.prediction_validation,
        )
        train_evaluator = EvalSurv(train_survival, durations_train, events_train, censor_surv="km")
        train_grid = evaluation_time_grid(
            durations_train,
            train_survival.index.values,
            self.config.evaluation_quantiles,
        )
        train_cindex = float(train_evaluator.concordance_td())
        train_ibs = float(train_evaluator.integrated_brier_score(train_grid))

        test_survival, prediction_diagnostics = self.adapter.predict_survival(
            model,
            data.test_data,
            validation_mode=self.config.prediction_validation,
        )
        test_evaluator = EvalSurv(test_survival, durations_test, events_test, censor_surv="km")
        test_grid = evaluation_time_grid(
            durations_test,
            test_survival.index.values,
            self.config.evaluation_quantiles,
        )
        test_cindex = float(test_evaluator.concordance_td())
        test_ibs = float(test_evaluator.integrated_brier_score(test_grid))

        self.server_ = server
        self.model_ = model
        self.labtrans_ = labtrans
        self.training_features_ = pooled_x
        self.training_labels_ = pooled_y
        self.last_prediction_diagnostics_ = prediction_diagnostics
        self.webdisco_fit_result_ = fit_result
        metadata = self.protocol.metadata()
        metadata.update(
            {
                "converged": bool(fit_result.converged),
                "n_iterations": int(fit_result.n_iterations),
                "score_norm": float(fit_result.score_norm),
                "payload_bytes_estimate": int(fit_result.communication_bytes),
                "training_prediction_diagnostics": {
                    "valid": bool(train_diagnostics.valid),
                    "identical_across_samples": bool(train_diagnostics.identical_across_samples),
                },
            }
        )
        return {
            "train_Cindex": [train_cindex],
            "train_IBS": [train_ibs],
            "test_Cindex": [test_cindex],
            "test_IBS": [test_ibs],
            "train_loss": [-float(fit_result.log_likelihood_history[-1])],
            "update_direction_norm": [float(np.linalg.norm(fit_result.coefficients))],
            "client_drift": [float("nan")],
            "update_direction_dispersion": [float("nan")],
            "empirical_heterogeneity": [float("nan")],
            "selected_clients": [[str(key) for key in clients_set]],
            "communication_bytes": [int(fit_result.communication_bytes)],
            "federated_protocol": self.protocol.name,
            "protocol_metadata": metadata,
            "telemetry_is_path_average_gradient": False,
            "telemetry_is_average_gradient": False,
            "telemetry_is_common_point_gradient": False,
            "telemetry_is_common_point_full_gradient": False,
        }

    def predict_survival(self, features: np.ndarray) -> pd.DataFrame:
        """Predict survival curves after :meth:`run` has fitted a model."""
        if not hasattr(self, "model_"):
            raise RuntimeError("run must be called before predict_survival")
        survival, diagnostics = self.adapter.predict_survival(
            self.model_,
            features,
            validation_mode=self.config.prediction_validation,
        )
        self.last_prediction_diagnostics_ = diagnostics
        return survival

    def get_privacy_info(self) -> Dict[str, Any]:
        """
        Get differential privacy protection information.

        Returns:
            dict: Privacy protection related information.
        """
        if not self.config.use_differential_privacy:
            return {"privacy_protection": False}

        from .differential_privacy import DifferentialPrivacy

        dp_tool = DifferentialPrivacy(self.config)

        # Compute privacy budget consumption
        total_epsilon, per_round_epsilon = dp_tool.compute_privacy_budget(
            self.config.global_epochs, self.config.num_clients
        )

        mechanism = self.config.dp_mechanism if hasattr(self.config, "dp_mechanism") else "gaussian"

        privacy_info = {
            "privacy_protection": True,
            "mechanism": mechanism,
            "epsilon": self.config.dp_epsilon,
            "sensitivity": self.config.dp_sensitivity,
            "total_epsilon": total_epsilon,
            "per_round_epsilon": per_round_epsilon,
            "formal_accounting_available": False,
            "privacy_scope": "experimental clipped client-update perturbation",
            # Which configured field actually scales the noise, and which
            # configured fields this mechanism ignores.  Reporting both keeps
            # the reported epsilon from being read as the noise control.
            "noise_driver": dp_tool.noise_driver(),
            "inactive_parameters": list(dp_tool.inactive_parameters()),
        }

        # Add Gaussian mechanism specific parameters
        if mechanism == "gaussian":
            privacy_info.update(
                {
                    "delta": self.config.dp_delta,
                    "noise_multiplier": self.config.dp_noise_multiplier,
                    "noise_scale": dp_tool.get_noise_scale(self.config.num_clients),
                    # Per-element standard deviation actually added to the
                    # clipped client update; compare it with typical weight
                    # magnitudes to tell whether the noise swamps the signal.
                    "effective_noise_sigma": (
                        dp_tool.get_noise_scale(self.config.num_clients)
                        * self.config.dp_clip_norm
                    ),
                }
            )

        # Add gradient clipping parameters (Gaussian and Laplace mechanisms)
        if mechanism in ["gaussian", "laplace"]:
            privacy_info["clip_norm"] = self.config.dp_clip_norm

        return privacy_info

    def _get_target(self, df):
        """Get the target variables."""
        return df[:, 0], df[:, 1]

    def plot_results(
        self,
        results: Dict[str, List[float]],
        output_path: Optional[Union[str, Path]] = None,
        show: bool = True,
    ) -> Figure:
        """
        Plot training results.

        Args:
            results: Training results in the form:
                {
                    'train_Cindex': List[float],  # Training set C-index
                    'train_IBS': List[float],     # Training set IBS
                    'test_Cindex': List[float],   # Test set C-index
                    'test_IBS': List[float]       # Test set IBS
                }
            output_path: Optional path to save the figure.
            show: Whether to open an interactive window; set to False for automated tests.

        Returns:
            matplotlib.figure.Figure: The generated figure object.
        """
        # Use a more modern style and avoid the deprecated seaborn style
        try:
            # Try the seaborn-v0_8 style (matplotlib 3.6+)
            plt.style.use("seaborn-v0_8")
        except OSError:
            # Fall back to the default style and set parameters manually
            plt.style.use("default")

        # Set fonts and style
        plt.rcParams.update(
            {
                "font.size": 12,
                "axes.labelsize": 14,
                "axes.titlesize": 16,
                "xtick.labelsize": 12,
                "ytick.labelsize": 12,
                "legend.fontsize": 12,
                "axes.spines.top": False,
                "axes.spines.right": False,
            }
        )

        # Close all existing figure windows to avoid multiple windows
        plt.close("all")

        # Create subplots
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))

        # Get the number of epochs and convert to a list
        n_epochs = len(results["train_Cindex"])
        epochs = list(range(1, n_epochs + 1))

        # Set the C-index axis range
        ax1.set_xlim(1, n_epochs)
        cindex_min = min(min(results["train_Cindex"]), min(results["test_Cindex"]))
        cindex_max = max(max(results["train_Cindex"]), max(results["test_Cindex"]))
        ax1.set_ylim(cindex_min - 0.05, cindex_max + 0.05)

        # Plot the C-index
        ax1.plot(
            epochs,
            results["train_Cindex"],
            color="#2ecc71",
            label="Training",
            linewidth=2,
            marker="o",
            markersize=4,
            markevery=5,
        )
        ax1.plot(
            epochs,
            results["test_Cindex"],
            color="#3498db",
            label="Testing",
            linewidth=2,
            linestyle="--",
            marker="s",
            markersize=4,
            markevery=5,
        )

        ax1.set_title("Concordance Index")
        ax1.set_xlabel("Epoch")
        ax1.set_ylabel("C-index")
        ax1.grid(True, linestyle="--", alpha=0.7)
        ax1.legend(loc="lower right", frameon=True, fancybox=True, shadow=True)

        # Add final value annotations
        final_train_cindex = results["train_Cindex"][-1]
        final_test_cindex = results["test_Cindex"][-1]
        ax1.text(
            0.02,
            0.98,
            f"Final Train C-index: {final_train_cindex:.3f}\nFinal Test C-index: {final_test_cindex:.3f}",
            transform=ax1.transAxes,
            bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"),
        )

        # Set the IBS axis range
        ax2.set_xlim(1, n_epochs)
        ibs_min = min(min(results["train_IBS"]), min(results["test_IBS"]))
        ibs_max = max(max(results["train_IBS"]), max(results["test_IBS"]))
        ax2.set_ylim(max(0, ibs_min - 0.05), min(1.0, ibs_max + 0.05))

        # Plot the IBS
        ax2.plot(
            epochs,
            results["train_IBS"],
            color="#2ecc71",
            label="Training",
            linewidth=2,
            marker="o",
            markersize=4,
            markevery=5,
        )
        ax2.plot(
            epochs,
            results["test_IBS"],
            color="#3498db",
            label="Testing",
            linewidth=2,
            linestyle="--",
            marker="s",
            markersize=4,
            markevery=5,
        )

        ax2.set_title("Integrated Brier Score")
        ax2.set_xlabel("Epoch")
        ax2.set_ylabel("IBS")
        ax2.grid(True, linestyle="--", alpha=0.7)
        ax2.legend(loc="upper right", frameon=True, fancybox=True, shadow=True)

        # Add final value annotations
        final_train_ibs = results["train_IBS"][-1]
        final_test_ibs = results["test_IBS"][-1]
        ax2.text(
            0.02,
            0.98,
            f"Final Train IBS: {final_train_ibs:.3f}\nFinal Test IBS: {final_test_ibs:.3f}",
            transform=ax2.transAxes,
            verticalalignment="top",
            bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"),
        )

        # Adjust layout
        plt.tight_layout()

        if output_path is not None:
            fig.savefig(output_path, dpi=300, bbox_inches="tight")
        if show:
            plt.show()
        return fig

__init__(config)

Initialize the federated survival analysis runner.

Parameters:

Name	Type	Description	Default
config	FSAConfig	Federated learning configuration.	required

Source code in federated_survival/core/runner.py

def __init__(self, config: FSAConfig):
    """
    Initialize the federated survival analysis runner.

    Args:
        config: Federated learning configuration.
    """
    self.config = config
    self.adapter = get_model_adapter(config.model_type)
    self.protocol = get_federated_protocol(config.federated_protocol)
    self.protocol.validate_config(config)
    self.set_random_seed()

set_random_seed()

Set the random seed.

Source code in federated_survival/core/runner.py

def set_random_seed(self):
    """Set the random seed."""
    if self.config.random_seed is not None:
        random.seed(self.config.random_seed)
        np.random.seed(self.config.random_seed)
        torch.manual_seed(self.config.random_seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(self.config.random_seed)

run(data, type='raw', aug_method='MVAEC')

Run federated survival analysis.

Parameters:

Name	Type	Description	Default
data	Any	Survival data with the following attributes: - clients_set: Client datasets in the form {client_id: (train_X, train_y)} - raw_aug_clients_set: Client datasets augmented from raw data, in the form {client_id: (train_X, train_y)} - test_data: Test data - test_label: Test labels	required
type	str	Data type, 'raw' or 'raw_aug'	'raw'
aug_method	str	Data augmentation method, 'MVAEC' or 'MVAES'	'MVAEC'

Returns:

Name	Type	Description
dict	Dict[str, List[float]]	{'train_Cindex': train_Cindex, 'train_IBS': train_IBS, 'test_Cindex': test_Cindex, 'test_IBS': test_IBS}

Source code in federated_survival/core/runner.py

def run(
    self, data: Any, type: str = "raw", aug_method: str = "MVAEC"
) -> Dict[str, List[float]]:
    """
    Run federated survival analysis.

    Args:
        data: Survival data with the following attributes:
            - clients_set: Client datasets in the form {client_id: (train_X, train_y)}
            - raw_aug_clients_set: Client datasets augmented from raw data,
              in the form {client_id: (train_X, train_y)}
            - test_data: Test data
            - test_label: Test labels
        type: Data type, 'raw' or 'raw_aug'
        aug_method: Data augmentation method, 'MVAEC' or 'MVAES'

    Returns:
        dict: {'train_Cindex': train_Cindex, 'train_IBS': train_IBS, 'test_Cindex': test_Cindex, 'test_IBS': test_IBS}
    """
    self.adapter = get_model_adapter(self.config.model_type)

    # Get training data
    if type == "raw":
        clients_set = data.clients_set
    elif type == "raw_aug":
        augmenter = DataAugmenter(
            latent_num=self.config.latent_num,
            hidden_num=self.config.hidden_num,
            alpha=self.config.alpha,
            beta=self.config.beta,
            sampling=self.config.augmentation_sampling,
            sparse_gamma=self.config.augmentation_sparse_gamma,
        )
        if aug_method == "MVAEC":
            clients_set = augmenter.mvaec(data.clients_set, k=self.config.k)
        elif aug_method == "MVAES":
            clients_set = augmenter.mvaes(data.clients_set, k=self.config.k)
        else:
            raise ValueError(f"Unsupported augmentation method: {aug_method}")
    else:
        raise ValueError(f"Unsupported data type: {type}")

    if self.protocol.owns_training_loop:
        if type != "raw":
            raise ValueError(
                f"protocol {self.protocol.name!r} does not support data augmentation"
            )
        return self._run_webdisco_style(data, clients_set)

    _y_train = np.empty(shape=(0, 2), dtype=np.float32)
    for i in clients_set.keys():
        _y_train = np.vstack((_y_train, clients_set[i][1]))
    y_train = self._get_target(_y_train)

    # Let the registered model adapter own label fitting and output shape.
    labtrans = self.adapter.configure_targets(self.config, *y_train)

    # Get test data
    durations_test, events_test = self._get_target(data.test_label)

    # Initialize server and clients
    server = Server(self.config)
    self.protocol.begin_run(self.config, server.global_model.state_dict())
    clients = []
    for i in clients_set.keys():
        clients.append(
            Client(
                self.config,
                server.global_model,
                clients_set[i],
                i,
                protocol=self.protocol,
            )
        )

    # If differential privacy is enabled, print privacy protection info
    if self.config.use_differential_privacy:
        if self.config.verbose:
            mechanism = (
                self.config.dp_mechanism if hasattr(self.config, "dp_mechanism") else "gaussian"
            )
            print("Differential privacy enabled:")
            print(f"  - mechanism: {mechanism.upper()}")
            print(f"  - privacy budget (epsilon): {self.config.dp_epsilon}")
            if mechanism == "gaussian":
                print(f"  - failure probability (delta): {self.config.dp_delta}")
                print(f"  - noise multiplier: {self.config.dp_noise_multiplier}")
            print(f"  - sensitivity: {self.config.dp_sensitivity}")
            if mechanism in ["gaussian", "laplace"]:
                print(f"  - gradient clipping norm: {self.config.dp_clip_norm}")

    # Record metrics
    train_Cindex = []
    train_IBS = []
    test_Cindex = []
    test_IBS = []

    # Federated learning training
    train_loss = []
    update_direction_norm = []
    client_drift = []
    update_direction_dispersion = []
    selected_clients = []
    communication_bytes = []
    model_bytes = sum(
        value.numel() * value.element_size()
        for value in server.global_model.state_dict().values()
    )

    for e in tqdm(
        range(self.config.global_epochs),
        disable=not self.config.show_progress,
    ):
        # Select clients
        candidates = random.sample(
            clients, max(round(self.config.client_sample_ratio * len(clients)), 1)
        )

        # Compute the total number of samples
        N = sum(c.N for c in candidates)
        round_directions = []
        round_weights = []
        round_losses = []
        round_drifts = []
        round_states = []
        global_state = {
            name: value.detach().clone()
            for name, value in server.global_model.state_dict().items()
        }

        # Client training
        for i, c in enumerate(candidates, 1):
            if self.config.verbose:
                print(f"Epoch {e+1}/{self.config.global_epochs}, Client {i}/{len(candidates)}")

            local = c.local_train(server.global_model, e)
            round_directions.append(c.last_update_direction)
            round_weights.append(c.N / N)
            round_losses.append(c.last_loss)
            round_drifts.append(c.last_drift_norm)
            round_states.append(
                {name: value.detach().clone() for name, value in local.state_dict().items()}
            )

        # Server update
        aggregated_state = self.protocol.aggregate(
            global_state,
            round_states,
            round_weights,
            round_index=e,
        )
        server.model_update(aggregated_state, len(candidates))

        weights = torch.tensor(round_weights, dtype=torch.float32)
        directions = torch.stack(round_directions)
        mean_direction = torch.sum(weights[:, None] * directions, dim=0)
        heterogeneity = torch.sum(
            weights * torch.sum((directions - mean_direction[None, :]) ** 2, dim=1)
        )
        train_loss.append(float(np.average(round_losses, weights=round_weights)))
        update_direction_norm.append(float(torch.linalg.vector_norm(mean_direction)))
        client_drift.append(float(np.average(round_drifts, weights=round_weights)))
        update_direction_dispersion.append(float(heterogeneity))
        selected_clients.append([str(c.client_id) for c in candidates])
        communication_bytes.append(
            self.protocol.round_communication_bytes(len(candidates), model_bytes)
        )

        # Evaluate the model
        index = server.model_eval(clients_set, labtrans)
        train_Cindex.append(index[0])
        train_IBS.append(index[1])

        # Compute test metrics
        model = self._get_model(server.global_model, labtrans)

        pooled_x = np.concatenate(
            [client_data[0] for client_data in clients_set.values()], axis=0
        )
        pooled_y = np.concatenate(
            [client_data[1] for client_data in clients_set.values()], axis=0
        )
        self.adapter.prepare_for_prediction(model, pooled_x, pooled_y, labtrans)
        surv, prediction_diagnostics = self.adapter.predict_survival(
            model,
            data.test_data,
            validation_mode=self.config.prediction_validation,
        )
        ev = EvalSurv(surv, durations_test, events_test, censor_surv="km")
        time_grid = evaluation_time_grid(
            durations_test,
            surv.index.values,
            self.config.evaluation_quantiles,
        )
        test_Cindex.append(ev.concordance_td())
        test_IBS.append(ev.integrated_brier_score(time_grid))

        if self.config.early_stopping:
            # Early stopping
            patience = self.config.early_stopping_patience
            if (
                len(train_Cindex) > patience
                and max(train_Cindex[-patience:]) <= train_Cindex[-patience - 1]
            ):
                break

    self.server_ = server
    self.model_ = model
    self.labtrans_ = labtrans
    self.training_features_ = pooled_x
    self.training_labels_ = pooled_y
    self.last_prediction_diagnostics_ = prediction_diagnostics

    return {
        "train_Cindex": train_Cindex,
        "train_IBS": train_IBS,
        "test_Cindex": test_Cindex,
        "test_IBS": test_IBS,
        "train_loss": train_loss,
        "update_direction_norm": update_direction_norm,
        "client_drift": client_drift,
        "update_direction_dispersion": update_direction_dispersion,
        # Backward-compatible alias. This is not the common-point B^2 in
        # the convergence assumption when local_epochs > 1.
        "empirical_heterogeneity": update_direction_dispersion,
        "selected_clients": selected_clients,
        "communication_bytes": communication_bytes,
        "federated_protocol": self.protocol.name,
        "protocol_metadata": self.protocol.metadata(),
        "telemetry_is_path_average_gradient": (
            self.config.optimizer == "sgd"
            and self.config.weight_decay == 0
            and not self.config.use_differential_privacy
        ),
        "telemetry_is_average_gradient": (
            self.config.optimizer == "sgd"
            and self.config.weight_decay == 0
            and not self.config.use_differential_privacy
        ),
        # With one local SGD step, every recorded direction is evaluated at
        # the broadcast point.  It equals the deterministic local-objective
        # gradient only in full-batch mode; otherwise it is a stochastic
        # gradient at that common point.
        "telemetry_is_common_point_gradient": (
            self.config.optimizer == "sgd"
            and self.config.weight_decay == 0
            and not self.config.use_differential_privacy
            and self.config.local_epochs == 1
        ),
        "telemetry_is_common_point_full_gradient": (
            self.config.optimizer == "sgd"
            and self.config.weight_decay == 0
            and not self.config.use_differential_privacy
            and self.config.local_epochs == 1
            and self.config.full_batch
        ),
    }

predict_survival(features)

Predict survival curves after :meth:run has fitted a model.

Source code in federated_survival/core/runner.py

def predict_survival(self, features: np.ndarray) -> pd.DataFrame:
    """Predict survival curves after :meth:`run` has fitted a model."""
    if not hasattr(self, "model_"):
        raise RuntimeError("run must be called before predict_survival")
    survival, diagnostics = self.adapter.predict_survival(
        self.model_,
        features,
        validation_mode=self.config.prediction_validation,
    )
    self.last_prediction_diagnostics_ = diagnostics
    return survival

get_privacy_info()

Get differential privacy protection information.

Returns:

Name	Type	Description
dict	Dict[str, Any]	Privacy protection related information.

Source code in federated_survival/core/runner.py

def get_privacy_info(self) -> Dict[str, Any]:
    """
    Get differential privacy protection information.

    Returns:
        dict: Privacy protection related information.
    """
    if not self.config.use_differential_privacy:
        return {"privacy_protection": False}

    from .differential_privacy import DifferentialPrivacy

    dp_tool = DifferentialPrivacy(self.config)

    # Compute privacy budget consumption
    total_epsilon, per_round_epsilon = dp_tool.compute_privacy_budget(
        self.config.global_epochs, self.config.num_clients
    )

    mechanism = self.config.dp_mechanism if hasattr(self.config, "dp_mechanism") else "gaussian"

    privacy_info = {
        "privacy_protection": True,
        "mechanism": mechanism,
        "epsilon": self.config.dp_epsilon,
        "sensitivity": self.config.dp_sensitivity,
        "total_epsilon": total_epsilon,
        "per_round_epsilon": per_round_epsilon,
        "formal_accounting_available": False,
        "privacy_scope": "experimental clipped client-update perturbation",
        # Which configured field actually scales the noise, and which
        # configured fields this mechanism ignores.  Reporting both keeps
        # the reported epsilon from being read as the noise control.
        "noise_driver": dp_tool.noise_driver(),
        "inactive_parameters": list(dp_tool.inactive_parameters()),
    }

    # Add Gaussian mechanism specific parameters
    if mechanism == "gaussian":
        privacy_info.update(
            {
                "delta": self.config.dp_delta,
                "noise_multiplier": self.config.dp_noise_multiplier,
                "noise_scale": dp_tool.get_noise_scale(self.config.num_clients),
                # Per-element standard deviation actually added to the
                # clipped client update; compare it with typical weight
                # magnitudes to tell whether the noise swamps the signal.
                "effective_noise_sigma": (
                    dp_tool.get_noise_scale(self.config.num_clients)
                    * self.config.dp_clip_norm
                ),
            }
        )

    # Add gradient clipping parameters (Gaussian and Laplace mechanisms)
    if mechanism in ["gaussian", "laplace"]:
        privacy_info["clip_norm"] = self.config.dp_clip_norm

    return privacy_info

plot_results(results, output_path=None, show=True)

Plot training results.

Parameters:

Name	Type	Description	Default
results	Dict[str, List[float]]	Training results in the form: { 'train_Cindex': List[float], # Training set C-index 'train_IBS': List[float], # Training set IBS 'test_Cindex': List[float], # Test set C-index 'test_IBS': List[float] # Test set IBS }	required
output_path	Optional[Union[str, Path]]	Optional path to save the figure.	None
show	bool	Whether to open an interactive window; set to False for automated tests.	True

Returns:

Type	Description
Figure	matplotlib.figure.Figure: The generated figure object.

Source code in federated_survival/core/runner.py

def plot_results(
    self,
    results: Dict[str, List[float]],
    output_path: Optional[Union[str, Path]] = None,
    show: bool = True,
) -> Figure:
    """
    Plot training results.

    Args:
        results: Training results in the form:
            {
                'train_Cindex': List[float],  # Training set C-index
                'train_IBS': List[float],     # Training set IBS
                'test_Cindex': List[float],   # Test set C-index
                'test_IBS': List[float]       # Test set IBS
            }
        output_path: Optional path to save the figure.
        show: Whether to open an interactive window; set to False for automated tests.

    Returns:
        matplotlib.figure.Figure: The generated figure object.
    """
    # Use a more modern style and avoid the deprecated seaborn style
    try:
        # Try the seaborn-v0_8 style (matplotlib 3.6+)
        plt.style.use("seaborn-v0_8")
    except OSError:
        # Fall back to the default style and set parameters manually
        plt.style.use("default")

    # Set fonts and style
    plt.rcParams.update(
        {
            "font.size": 12,
            "axes.labelsize": 14,
            "axes.titlesize": 16,
            "xtick.labelsize": 12,
            "ytick.labelsize": 12,
            "legend.fontsize": 12,
            "axes.spines.top": False,
            "axes.spines.right": False,
        }
    )

    # Close all existing figure windows to avoid multiple windows
    plt.close("all")

    # Create subplots
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))

    # Get the number of epochs and convert to a list
    n_epochs = len(results["train_Cindex"])
    epochs = list(range(1, n_epochs + 1))

    # Set the C-index axis range
    ax1.set_xlim(1, n_epochs)
    cindex_min = min(min(results["train_Cindex"]), min(results["test_Cindex"]))
    cindex_max = max(max(results["train_Cindex"]), max(results["test_Cindex"]))
    ax1.set_ylim(cindex_min - 0.05, cindex_max + 0.05)

    # Plot the C-index
    ax1.plot(
        epochs,
        results["train_Cindex"],
        color="#2ecc71",
        label="Training",
        linewidth=2,
        marker="o",
        markersize=4,
        markevery=5,
    )
    ax1.plot(
        epochs,
        results["test_Cindex"],
        color="#3498db",
        label="Testing",
        linewidth=2,
        linestyle="--",
        marker="s",
        markersize=4,
        markevery=5,
    )

    ax1.set_title("Concordance Index")
    ax1.set_xlabel("Epoch")
    ax1.set_ylabel("C-index")
    ax1.grid(True, linestyle="--", alpha=0.7)
    ax1.legend(loc="lower right", frameon=True, fancybox=True, shadow=True)

    # Add final value annotations
    final_train_cindex = results["train_Cindex"][-1]
    final_test_cindex = results["test_Cindex"][-1]
    ax1.text(
        0.02,
        0.98,
        f"Final Train C-index: {final_train_cindex:.3f}\nFinal Test C-index: {final_test_cindex:.3f}",
        transform=ax1.transAxes,
        bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"),
    )

    # Set the IBS axis range
    ax2.set_xlim(1, n_epochs)
    ibs_min = min(min(results["train_IBS"]), min(results["test_IBS"]))
    ibs_max = max(max(results["train_IBS"]), max(results["test_IBS"]))
    ax2.set_ylim(max(0, ibs_min - 0.05), min(1.0, ibs_max + 0.05))

    # Plot the IBS
    ax2.plot(
        epochs,
        results["train_IBS"],
        color="#2ecc71",
        label="Training",
        linewidth=2,
        marker="o",
        markersize=4,
        markevery=5,
    )
    ax2.plot(
        epochs,
        results["test_IBS"],
        color="#3498db",
        label="Testing",
        linewidth=2,
        linestyle="--",
        marker="s",
        markersize=4,
        markevery=5,
    )

    ax2.set_title("Integrated Brier Score")
    ax2.set_xlabel("Epoch")
    ax2.set_ylabel("IBS")
    ax2.grid(True, linestyle="--", alpha=0.7)
    ax2.legend(loc="upper right", frameon=True, fancybox=True, shadow=True)

    # Add final value annotations
    final_train_ibs = results["train_IBS"][-1]
    final_test_ibs = results["test_IBS"][-1]
    ax2.text(
        0.02,
        0.98,
        f"Final Train IBS: {final_train_ibs:.3f}\nFinal Test IBS: {final_test_ibs:.3f}",
        transform=ax2.transAxes,
        verticalalignment="top",
        bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"),
    )

    # Adjust layout
    plt.tight_layout()

    if output_path is not None:
        fig.savefig(output_path, dpi=300, bbox_inches="tight")
    if show:
        plt.show()
    return fig