826 lines
24 KiB
Rust
826 lines
24 KiB
Rust
//! Shared types for FederatedMed - Privacy-Preserving Medical AI.
|
|
//!
|
|
//! This crate defines the IPC types for federated learning on hospital data
|
|
//! without sharing sensitive patient information.
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
// ============================================================================
|
|
// Federated Learning Configuration
|
|
// ============================================================================
|
|
|
|
/// Configuration for federated learning training.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct FederatedConfig {
|
|
/// Number of federated rounds.
|
|
pub num_rounds: usize,
|
|
/// Number of local epochs per round.
|
|
pub local_epochs: usize,
|
|
/// Batch size for local training.
|
|
pub batch_size: usize,
|
|
/// Aggregation strategy.
|
|
pub aggregation_strategy: AggregationStrategy,
|
|
/// Learning rate.
|
|
pub learning_rate: f64,
|
|
/// Minimum number of clients per round.
|
|
pub min_clients: usize,
|
|
/// Client selection fraction (0.0 to 1.0).
|
|
pub client_fraction: f32,
|
|
/// Whether to use secure aggregation.
|
|
pub secure_aggregation: bool,
|
|
/// Privacy configuration.
|
|
pub privacy: Option<PrivacyConfig>,
|
|
}
|
|
|
|
impl Default for FederatedConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
num_rounds: 100,
|
|
local_epochs: 5,
|
|
batch_size: 32,
|
|
aggregation_strategy: AggregationStrategy::FedAvg,
|
|
learning_rate: 0.01,
|
|
min_clients: 3,
|
|
client_fraction: 0.3,
|
|
secure_aggregation: true,
|
|
privacy: Some(PrivacyConfig::default()),
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Aggregation Strategy
|
|
// ============================================================================
|
|
|
|
/// Aggregation strategy for combining client updates.
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
|
|
pub enum AggregationStrategy {
|
|
/// Federated Averaging - weighted average of model updates.
|
|
FedAvg,
|
|
/// FedProx - adds proximal term for heterogeneous data.
|
|
FedProx,
|
|
/// Scaffold - uses control variates to correct client drift.
|
|
Scaffold,
|
|
/// FedAdam - adaptive learning rate aggregation.
|
|
FedAdam,
|
|
/// FedYogi - adaptive aggregation with momentum.
|
|
FedYogi,
|
|
/// TrimmedMean - robust aggregation against Byzantine clients.
|
|
TrimmedMean,
|
|
/// Median - median-based robust aggregation.
|
|
Median,
|
|
/// Krum - Byzantine-resilient aggregation.
|
|
Krum,
|
|
}
|
|
|
|
/// Parameters specific to FedProx.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct FedProxParams {
|
|
/// Proximal term coefficient (mu).
|
|
pub mu: f64,
|
|
}
|
|
|
|
impl Default for FedProxParams {
|
|
fn default() -> Self {
|
|
Self { mu: 0.01 }
|
|
}
|
|
}
|
|
|
|
/// Parameters specific to Scaffold.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ScaffoldParams {
|
|
/// Server learning rate.
|
|
pub server_lr: f64,
|
|
/// Client learning rate.
|
|
pub client_lr: f64,
|
|
}
|
|
|
|
impl Default for ScaffoldParams {
|
|
fn default() -> Self {
|
|
Self {
|
|
server_lr: 1.0,
|
|
client_lr: 0.1,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Client Information
|
|
// ============================================================================
|
|
|
|
/// Information about a federated learning client (e.g., a hospital).
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ClientInfo {
|
|
/// Unique client identifier.
|
|
pub id: String,
|
|
/// Human-readable name (e.g., "General Hospital").
|
|
pub name: String,
|
|
/// Number of data samples at this client.
|
|
pub data_size: usize,
|
|
/// Whether the client is currently active.
|
|
pub is_active: bool,
|
|
/// Client type.
|
|
pub client_type: ClientType,
|
|
/// Data distribution characteristics.
|
|
pub data_distribution: DataDistribution,
|
|
/// Last communication timestamp (Unix epoch).
|
|
pub last_seen: u64,
|
|
/// Network latency in milliseconds.
|
|
pub latency_ms: u32,
|
|
}
|
|
|
|
impl Default for ClientInfo {
|
|
fn default() -> Self {
|
|
Self {
|
|
id: "client_0".to_string(),
|
|
name: "Default Hospital".to_string(),
|
|
data_size: 1000,
|
|
is_active: true,
|
|
client_type: ClientType::Hospital,
|
|
data_distribution: DataDistribution::default(),
|
|
last_seen: 0,
|
|
latency_ms: 50,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Type of federated learning client.
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
|
|
pub enum ClientType {
|
|
/// Hospital.
|
|
Hospital,
|
|
/// Clinic.
|
|
Clinic,
|
|
/// Research institution.
|
|
Research,
|
|
/// Radiology center.
|
|
Radiology,
|
|
/// Laboratory.
|
|
Laboratory,
|
|
}
|
|
|
|
/// Data distribution characteristics at a client.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct DataDistribution {
|
|
/// Class distribution (class_id -> count).
|
|
pub class_counts: Vec<usize>,
|
|
/// Whether data is IID (identically and independently distributed).
|
|
pub is_iid: bool,
|
|
/// Data heterogeneity measure (0.0 = homogeneous, 1.0 = highly heterogeneous).
|
|
pub heterogeneity: f32,
|
|
}
|
|
|
|
impl Default for DataDistribution {
|
|
fn default() -> Self {
|
|
Self {
|
|
class_counts: vec![500, 300, 200],
|
|
is_iid: false,
|
|
heterogeneity: 0.3,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Privacy Configuration
|
|
// ============================================================================
|
|
|
|
/// Privacy configuration for differential privacy.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct PrivacyConfig {
|
|
/// Epsilon - privacy budget (lower = more private).
|
|
pub epsilon: f64,
|
|
/// Delta - probability of privacy breach.
|
|
pub delta: f64,
|
|
/// Gradient clipping norm for sensitivity bounding.
|
|
pub clip_norm: f64,
|
|
/// Noise multiplier for Gaussian mechanism.
|
|
pub noise_multiplier: f64,
|
|
/// Whether to use local differential privacy.
|
|
pub local_dp: bool,
|
|
/// Accountant type for privacy budget tracking.
|
|
pub accountant: PrivacyAccountantType,
|
|
/// Target epsilon for automatic noise calibration.
|
|
pub target_epsilon: Option<f64>,
|
|
}
|
|
|
|
impl Default for PrivacyConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
epsilon: 8.0,
|
|
delta: 1e-5,
|
|
clip_norm: 1.0,
|
|
noise_multiplier: 1.1,
|
|
local_dp: false,
|
|
accountant: PrivacyAccountantType::RDP,
|
|
target_epsilon: Some(10.0),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Type of privacy accountant.
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
|
|
pub enum PrivacyAccountantType {
|
|
/// Moments Accountant.
|
|
Moments,
|
|
/// Renyi Differential Privacy.
|
|
RDP,
|
|
/// Gaussian Differential Privacy.
|
|
GDP,
|
|
/// Privacy Loss Distribution.
|
|
PLD,
|
|
}
|
|
|
|
// ============================================================================
|
|
// Model Types
|
|
// ============================================================================
|
|
|
|
/// Global model maintained by the server.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct GlobalModel {
|
|
/// Model version/round number.
|
|
pub version: u64,
|
|
/// Model weights (flattened).
|
|
pub weights: Vec<f32>,
|
|
/// Weight shapes for reconstruction.
|
|
pub shapes: Vec<Vec<usize>>,
|
|
/// Layer names.
|
|
pub layer_names: Vec<String>,
|
|
/// Model architecture type.
|
|
pub architecture: ModelArchitecture,
|
|
/// Number of parameters.
|
|
pub num_parameters: usize,
|
|
/// Accuracy on validation set.
|
|
pub validation_accuracy: Option<f64>,
|
|
/// Loss on validation set.
|
|
pub validation_loss: Option<f64>,
|
|
}
|
|
|
|
impl Default for GlobalModel {
|
|
fn default() -> Self {
|
|
// Create a simple model structure
|
|
let shapes = vec![
|
|
vec![784, 256], // Input to hidden1
|
|
vec![256], // Hidden1 bias
|
|
vec![256, 128], // Hidden1 to hidden2
|
|
vec![128], // Hidden2 bias
|
|
vec![128, 10], // Hidden2 to output
|
|
vec![10], // Output bias
|
|
];
|
|
let num_params: usize = shapes.iter().map(|s| s.iter().product::<usize>()).sum();
|
|
|
|
Self {
|
|
version: 0,
|
|
weights: vec![0.0; num_params],
|
|
shapes,
|
|
layer_names: vec![
|
|
"fc1.weight".to_string(),
|
|
"fc1.bias".to_string(),
|
|
"fc2.weight".to_string(),
|
|
"fc2.bias".to_string(),
|
|
"fc3.weight".to_string(),
|
|
"fc3.bias".to_string(),
|
|
],
|
|
architecture: ModelArchitecture::CNN,
|
|
num_parameters: num_params,
|
|
validation_accuracy: None,
|
|
validation_loss: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Model architecture type.
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
|
|
pub enum ModelArchitecture {
|
|
/// Convolutional neural network.
|
|
CNN,
|
|
/// ResNet variant.
|
|
ResNet,
|
|
/// DenseNet variant.
|
|
DenseNet,
|
|
/// Vision Transformer.
|
|
ViT,
|
|
/// U-Net for segmentation.
|
|
UNet,
|
|
/// Multi-layer perceptron.
|
|
MLP,
|
|
/// Transformer.
|
|
Transformer,
|
|
}
|
|
|
|
/// Local update from a client.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct LocalUpdate {
|
|
/// Client ID.
|
|
pub client_id: String,
|
|
/// Round number.
|
|
pub round: u64,
|
|
/// Weight updates (delta from global model).
|
|
pub weight_deltas: Vec<f32>,
|
|
/// Number of local samples used.
|
|
pub num_samples: usize,
|
|
/// Local training loss.
|
|
pub local_loss: f64,
|
|
/// Local training accuracy.
|
|
pub local_accuracy: f64,
|
|
/// Training time in seconds.
|
|
pub training_time: f64,
|
|
/// Control variates for Scaffold (optional).
|
|
pub control_variates: Option<Vec<f32>>,
|
|
}
|
|
|
|
impl Default for LocalUpdate {
|
|
fn default() -> Self {
|
|
Self {
|
|
client_id: "client_0".to_string(),
|
|
round: 0,
|
|
weight_deltas: vec![],
|
|
num_samples: 0,
|
|
local_loss: 0.0,
|
|
local_accuracy: 0.0,
|
|
training_time: 0.0,
|
|
control_variates: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Training Progress
|
|
// ============================================================================
|
|
|
|
/// Training progress for federated learning.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct TrainingProgress {
|
|
/// Current round.
|
|
pub round: u64,
|
|
/// Total rounds.
|
|
pub total_rounds: u64,
|
|
/// Global model accuracy.
|
|
pub accuracy: f64,
|
|
/// Global model loss.
|
|
pub loss: f64,
|
|
/// Privacy budget spent so far.
|
|
pub privacy_budget_spent: f64,
|
|
/// Number of participating clients this round.
|
|
pub participating_clients: usize,
|
|
/// Total number of clients.
|
|
pub total_clients: usize,
|
|
/// Elapsed time in seconds.
|
|
pub elapsed_seconds: f64,
|
|
/// Estimated time remaining in seconds.
|
|
pub eta_seconds: Option<f64>,
|
|
/// Current phase.
|
|
pub phase: TrainingPhase,
|
|
}
|
|
|
|
impl Default for TrainingProgress {
|
|
fn default() -> Self {
|
|
Self {
|
|
round: 0,
|
|
total_rounds: 100,
|
|
accuracy: 0.0,
|
|
loss: f64::INFINITY,
|
|
privacy_budget_spent: 0.0,
|
|
participating_clients: 0,
|
|
total_clients: 0,
|
|
elapsed_seconds: 0.0,
|
|
eta_seconds: None,
|
|
phase: TrainingPhase::Initialization,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Training phase.
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
|
|
pub enum TrainingPhase {
|
|
/// Initialization phase.
|
|
Initialization,
|
|
/// Client selection phase.
|
|
ClientSelection,
|
|
/// Local training phase.
|
|
LocalTraining,
|
|
/// Aggregation phase.
|
|
Aggregation,
|
|
/// Evaluation phase.
|
|
Evaluation,
|
|
/// Completed.
|
|
Completed,
|
|
}
|
|
|
|
// ============================================================================
|
|
// Medical Domain Types
|
|
// ============================================================================
|
|
|
|
/// Medical imaging task type.
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
|
|
pub enum MedicalTask {
|
|
/// Chest X-ray classification (e.g., COVID, pneumonia).
|
|
ChestXrayClassification,
|
|
/// Skin lesion classification (melanoma detection).
|
|
SkinLesionClassification,
|
|
/// Retinal disease detection.
|
|
RetinalDisease,
|
|
/// Brain MRI segmentation.
|
|
BrainMRISegmentation,
|
|
/// CT scan analysis.
|
|
CTScanAnalysis,
|
|
/// Mammography screening.
|
|
MammographyScreening,
|
|
/// Pathology slide classification.
|
|
PathologySlide,
|
|
/// ECG classification.
|
|
ECGClassification,
|
|
}
|
|
|
|
/// Medical dataset information.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct MedicalDataset {
|
|
/// Dataset name.
|
|
pub name: String,
|
|
/// Task type.
|
|
pub task: MedicalTask,
|
|
/// Number of samples.
|
|
pub num_samples: usize,
|
|
/// Number of classes.
|
|
pub num_classes: usize,
|
|
/// Class names.
|
|
pub class_names: Vec<String>,
|
|
/// Image dimensions (height, width, channels).
|
|
pub image_dims: (usize, usize, usize),
|
|
/// Whether dataset is labeled.
|
|
pub is_labeled: bool,
|
|
}
|
|
|
|
impl Default for MedicalDataset {
|
|
fn default() -> Self {
|
|
Self {
|
|
name: "ChestXray14".to_string(),
|
|
task: MedicalTask::ChestXrayClassification,
|
|
num_samples: 112120,
|
|
num_classes: 14,
|
|
class_names: vec![
|
|
"Atelectasis".to_string(),
|
|
"Cardiomegaly".to_string(),
|
|
"Effusion".to_string(),
|
|
"Infiltration".to_string(),
|
|
"Mass".to_string(),
|
|
"Nodule".to_string(),
|
|
"Pneumonia".to_string(),
|
|
"Pneumothorax".to_string(),
|
|
"Consolidation".to_string(),
|
|
"Edema".to_string(),
|
|
"Emphysema".to_string(),
|
|
"Fibrosis".to_string(),
|
|
"Pleural Thickening".to_string(),
|
|
"Hernia".to_string(),
|
|
],
|
|
image_dims: (224, 224, 1),
|
|
is_labeled: true,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Federation Results
|
|
// ============================================================================
|
|
|
|
/// Final result of federated training.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct FederationResult {
|
|
/// Final global model.
|
|
pub final_model: GlobalModel,
|
|
/// Training history (accuracy per round).
|
|
pub accuracy_history: Vec<f64>,
|
|
/// Training history (loss per round).
|
|
pub loss_history: Vec<f64>,
|
|
/// Total training time.
|
|
pub total_time: f64,
|
|
/// Number of rounds completed.
|
|
pub rounds_completed: u64,
|
|
/// Total privacy budget spent.
|
|
pub total_privacy_budget: f64,
|
|
/// Client participation statistics.
|
|
pub client_stats: Vec<ClientStats>,
|
|
/// Final test metrics.
|
|
pub test_metrics: TestMetrics,
|
|
}
|
|
|
|
/// Statistics for a single client.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ClientStats {
|
|
/// Client ID.
|
|
pub client_id: String,
|
|
/// Number of rounds participated.
|
|
pub rounds_participated: usize,
|
|
/// Average local accuracy.
|
|
pub avg_local_accuracy: f64,
|
|
/// Total samples contributed.
|
|
pub total_samples: usize,
|
|
/// Average training time per round.
|
|
pub avg_training_time: f64,
|
|
}
|
|
|
|
/// Test metrics for evaluation.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct TestMetrics {
|
|
/// Accuracy.
|
|
pub accuracy: f64,
|
|
/// Precision (macro-averaged).
|
|
pub precision: f64,
|
|
/// Recall (macro-averaged).
|
|
pub recall: f64,
|
|
/// F1 score (macro-averaged).
|
|
pub f1_score: f64,
|
|
/// AUC-ROC (macro-averaged).
|
|
pub auc_roc: f64,
|
|
/// Confusion matrix (flattened).
|
|
pub confusion_matrix: Vec<usize>,
|
|
/// Number of classes.
|
|
pub num_classes: usize,
|
|
}
|
|
|
|
impl Default for TestMetrics {
|
|
fn default() -> Self {
|
|
Self {
|
|
accuracy: 0.0,
|
|
precision: 0.0,
|
|
recall: 0.0,
|
|
f1_score: 0.0,
|
|
auc_roc: 0.0,
|
|
confusion_matrix: vec![],
|
|
num_classes: 0,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Sample Data Functions
|
|
// ============================================================================
|
|
|
|
/// Create a sample federated config.
|
|
#[must_use]
|
|
pub fn sample_federated_config() -> FederatedConfig {
|
|
FederatedConfig {
|
|
num_rounds: 50,
|
|
local_epochs: 3,
|
|
batch_size: 32,
|
|
aggregation_strategy: AggregationStrategy::FedAvg,
|
|
learning_rate: 0.01,
|
|
min_clients: 3,
|
|
client_fraction: 0.5,
|
|
secure_aggregation: true,
|
|
privacy: Some(PrivacyConfig {
|
|
epsilon: 8.0,
|
|
delta: 1e-5,
|
|
clip_norm: 1.0,
|
|
noise_multiplier: 1.1,
|
|
local_dp: false,
|
|
accountant: PrivacyAccountantType::RDP,
|
|
target_epsilon: Some(10.0),
|
|
}),
|
|
}
|
|
}
|
|
|
|
/// Create sample client infos.
|
|
#[must_use]
|
|
pub fn sample_clients() -> Vec<ClientInfo> {
|
|
vec![
|
|
ClientInfo {
|
|
id: "hospital_a".to_string(),
|
|
name: "City General Hospital".to_string(),
|
|
data_size: 5000,
|
|
is_active: true,
|
|
client_type: ClientType::Hospital,
|
|
data_distribution: DataDistribution {
|
|
class_counts: vec![1500, 1200, 800, 700, 500, 300],
|
|
is_iid: false,
|
|
heterogeneity: 0.25,
|
|
},
|
|
last_seen: 1704067200,
|
|
latency_ms: 45,
|
|
},
|
|
ClientInfo {
|
|
id: "hospital_b".to_string(),
|
|
name: "University Medical Center".to_string(),
|
|
data_size: 8000,
|
|
is_active: true,
|
|
client_type: ClientType::Hospital,
|
|
data_distribution: DataDistribution {
|
|
class_counts: vec![2000, 1800, 1500, 1200, 900, 600],
|
|
is_iid: false,
|
|
heterogeneity: 0.2,
|
|
},
|
|
last_seen: 1704067200,
|
|
latency_ms: 30,
|
|
},
|
|
ClientInfo {
|
|
id: "clinic_c".to_string(),
|
|
name: "Downtown Health Clinic".to_string(),
|
|
data_size: 2000,
|
|
is_active: true,
|
|
client_type: ClientType::Clinic,
|
|
data_distribution: DataDistribution {
|
|
class_counts: vec![600, 500, 400, 300, 150, 50],
|
|
is_iid: false,
|
|
heterogeneity: 0.4,
|
|
},
|
|
last_seen: 1704067200,
|
|
latency_ms: 60,
|
|
},
|
|
ClientInfo {
|
|
id: "research_d".to_string(),
|
|
name: "National Research Institute".to_string(),
|
|
data_size: 3500,
|
|
is_active: true,
|
|
client_type: ClientType::Research,
|
|
data_distribution: DataDistribution {
|
|
class_counts: vec![700, 700, 700, 700, 350, 350],
|
|
is_iid: true,
|
|
heterogeneity: 0.1,
|
|
},
|
|
last_seen: 1704067200,
|
|
latency_ms: 25,
|
|
},
|
|
ClientInfo {
|
|
id: "radiology_e".to_string(),
|
|
name: "Regional Radiology Center".to_string(),
|
|
data_size: 4500,
|
|
is_active: true,
|
|
client_type: ClientType::Radiology,
|
|
data_distribution: DataDistribution {
|
|
class_counts: vec![1200, 1000, 900, 800, 400, 200],
|
|
is_iid: false,
|
|
heterogeneity: 0.3,
|
|
},
|
|
last_seen: 1704067200,
|
|
latency_ms: 40,
|
|
},
|
|
]
|
|
}
|
|
|
|
/// Create a sample global model.
|
|
#[must_use]
|
|
pub fn sample_global_model() -> GlobalModel {
|
|
GlobalModel::default()
|
|
}
|
|
|
|
/// Create sample privacy config.
|
|
#[must_use]
|
|
pub fn sample_privacy_config() -> PrivacyConfig {
|
|
PrivacyConfig::default()
|
|
}
|
|
|
|
/// Create a sample training progress.
|
|
#[must_use]
|
|
pub fn sample_training_progress() -> TrainingProgress {
|
|
TrainingProgress {
|
|
round: 25,
|
|
total_rounds: 100,
|
|
accuracy: 0.85,
|
|
loss: 0.45,
|
|
privacy_budget_spent: 4.2,
|
|
participating_clients: 4,
|
|
total_clients: 5,
|
|
elapsed_seconds: 3600.0,
|
|
eta_seconds: Some(3600.0),
|
|
phase: TrainingPhase::Aggregation,
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Tests
|
|
// ============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_federated_config() {
|
|
let config = sample_federated_config();
|
|
assert_eq!(config.num_rounds, 50);
|
|
assert_eq!(config.aggregation_strategy, AggregationStrategy::FedAvg);
|
|
assert!(config.privacy.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_aggregation_strategy() {
|
|
let strategies = vec![
|
|
AggregationStrategy::FedAvg,
|
|
AggregationStrategy::FedProx,
|
|
AggregationStrategy::Scaffold,
|
|
];
|
|
for strategy in strategies {
|
|
let json = serde_json::to_string(&strategy).unwrap();
|
|
let parsed: AggregationStrategy = serde_json::from_str(&json).unwrap();
|
|
assert_eq!(parsed, strategy);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_client_info() {
|
|
let clients = sample_clients();
|
|
assert_eq!(clients.len(), 5);
|
|
assert!(clients.iter().all(|c| c.is_active));
|
|
}
|
|
|
|
#[test]
|
|
fn test_privacy_config() {
|
|
let config = sample_privacy_config();
|
|
assert!(config.epsilon > 0.0);
|
|
assert!(config.delta > 0.0);
|
|
assert!(config.clip_norm > 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_global_model() {
|
|
let model = sample_global_model();
|
|
assert_eq!(model.version, 0);
|
|
assert!(!model.weights.is_empty());
|
|
assert!(!model.shapes.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_local_update() {
|
|
let update = LocalUpdate {
|
|
client_id: "hospital_a".to_string(),
|
|
round: 10,
|
|
weight_deltas: vec![0.1, 0.2, 0.3],
|
|
num_samples: 1000,
|
|
local_loss: 0.5,
|
|
local_accuracy: 0.85,
|
|
training_time: 120.0,
|
|
control_variates: None,
|
|
};
|
|
assert_eq!(update.round, 10);
|
|
assert_eq!(update.num_samples, 1000);
|
|
}
|
|
|
|
#[test]
|
|
fn test_training_progress() {
|
|
let progress = sample_training_progress();
|
|
assert_eq!(progress.round, 25);
|
|
assert!(progress.accuracy > 0.0);
|
|
assert_eq!(progress.phase, TrainingPhase::Aggregation);
|
|
}
|
|
|
|
#[test]
|
|
fn test_medical_task() {
|
|
let tasks = vec![
|
|
MedicalTask::ChestXrayClassification,
|
|
MedicalTask::SkinLesionClassification,
|
|
MedicalTask::RetinalDisease,
|
|
];
|
|
for task in tasks {
|
|
let json = serde_json::to_string(&task).unwrap();
|
|
let parsed: MedicalTask = serde_json::from_str(&json).unwrap();
|
|
assert_eq!(parsed, task);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_medical_dataset() {
|
|
let dataset = MedicalDataset::default();
|
|
assert_eq!(dataset.num_classes, 14);
|
|
assert_eq!(dataset.class_names.len(), 14);
|
|
}
|
|
|
|
#[test]
|
|
fn test_test_metrics() {
|
|
let metrics = TestMetrics {
|
|
accuracy: 0.92,
|
|
precision: 0.90,
|
|
recall: 0.88,
|
|
f1_score: 0.89,
|
|
auc_roc: 0.95,
|
|
confusion_matrix: vec![100, 5, 3, 92],
|
|
num_classes: 2,
|
|
};
|
|
assert!(metrics.accuracy > 0.9);
|
|
assert!(metrics.auc_roc > 0.9);
|
|
}
|
|
|
|
#[test]
|
|
fn test_serialization() {
|
|
let config = sample_federated_config();
|
|
let json = serde_json::to_string(&config).unwrap();
|
|
let parsed: FederatedConfig = serde_json::from_str(&json).unwrap();
|
|
assert_eq!(parsed.num_rounds, config.num_rounds);
|
|
}
|
|
|
|
#[test]
|
|
fn test_federation_result() {
|
|
let result = FederationResult {
|
|
final_model: sample_global_model(),
|
|
accuracy_history: vec![0.5, 0.6, 0.7, 0.8, 0.85],
|
|
loss_history: vec![2.0, 1.5, 1.0, 0.7, 0.5],
|
|
total_time: 7200.0,
|
|
rounds_completed: 50,
|
|
total_privacy_budget: 8.5,
|
|
client_stats: vec![],
|
|
test_metrics: TestMetrics::default(),
|
|
};
|
|
assert_eq!(result.accuracy_history.len(), 5);
|
|
assert_eq!(result.rounds_completed, 50);
|
|
}
|
|
}
|