idea141-openimpact-privacy-.../src/openimpact/evaluation.py

41 lines
1.6 KiB
Python

from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from .contracts import flatten_batches
from .federated import FederatedTrainer
from .model import TemporaryImpactModel
from .replay import DeterministicReplayEngine
from .synthetic import SyntheticMarket
@dataclass(frozen=True)
class EvaluationMetrics:
rmse: float
latency_ms: float
privacy_budget: float
leakage_bound: float
def _rmse(prediction: np.ndarray, target: np.ndarray) -> float:
return float(np.sqrt(np.mean((prediction - target) ** 2)))
def evaluate_federated_setup(market: SyntheticMarket, trainer: FederatedTrainer | None = None) -> tuple[EvaluationMetrics, TemporaryImpactModel, TemporaryImpactModel]:
trainer = trainer or FederatedTrainer()
flat_requests, flat_impacts = flatten_batches(market.batches)
baseline = TemporaryImpactModel().fit(flat_requests, flat_impacts)
result = trainer.fit(market.batches)
replay = DeterministicReplayEngine(seed=trainer.seed)
steps = replay.replay(flat_requests, result.global_model, flat_impacts)
latency_ms = float(np.mean([(step.delivery_ns - step.timestamp_ns) / 1_000_000.0 for step in steps]))
predictions = result.global_model.predict(flat_requests)
rmse = _rmse(predictions, flat_impacts)
privacy_budget = float(sum(update.privacy_budget for update in result.updates))
leakage_bound = float(np.max(np.abs(predictions - baseline.predict(flat_requests))))
return EvaluationMetrics(rmse=rmse, latency_ms=latency_ms, privacy_budget=privacy_budget, leakage_bound=leakage_bound), result.global_model, baseline