54 lines
2.2 KiB
Python
54 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
from openimpact.evaluation import evaluate_federated_setup
|
|
from openimpact.federated import FederatedTrainer
|
|
from openimpact.ledger import GovernanceLedger
|
|
from openimpact.model import TemporaryImpactModel
|
|
from openimpact.replay import DeterministicReplayEngine
|
|
from openimpact.synthetic import SyntheticMarketConfig, generate_synthetic_market
|
|
|
|
|
|
def test_temporary_impact_model_fits_synthetic_market():
|
|
market = generate_synthetic_market(SyntheticMarketConfig(seed=11, venue_count=2, samples_per_venue=64))
|
|
requests = market.requests
|
|
impacts = market.impacts
|
|
|
|
model = TemporaryImpactModel().fit(requests, impacts)
|
|
predictions = model.predict(requests)
|
|
|
|
assert predictions.shape == impacts.shape
|
|
assert float(np.sqrt(np.mean((predictions - impacts) ** 2))) < 0.05
|
|
|
|
|
|
def test_federated_training_and_evaluation_are_deterministic():
|
|
market = generate_synthetic_market(SyntheticMarketConfig(seed=3, venue_count=3, samples_per_venue=32))
|
|
trainer = FederatedTrainer(seed=3)
|
|
metrics_a, model_a, baseline_a = evaluate_federated_setup(market, trainer)
|
|
metrics_b, model_b, baseline_b = evaluate_federated_setup(market, trainer)
|
|
|
|
assert metrics_a == metrics_b
|
|
assert np.allclose(model_a.coefficients_, model_b.coefficients_)
|
|
assert np.allclose(baseline_a.coefficients_, baseline_b.coefficients_)
|
|
assert metrics_a.rmse < 0.08
|
|
|
|
|
|
def test_replay_and_ledger_work_end_to_end(tmp_path: Path):
|
|
market = generate_synthetic_market(SyntheticMarketConfig(seed=5, venue_count=2, samples_per_venue=16))
|
|
trainer = FederatedTrainer(seed=5)
|
|
result = trainer.fit(market.batches)
|
|
replay = DeterministicReplayEngine(seed=5, base_latency_ms={"venue-1": 2.5, "venue-2": 3.5})
|
|
steps = replay.replay(market.requests, result.global_model, market.impacts)
|
|
|
|
ledger = GovernanceLedger(tmp_path / "ledger.sqlite3")
|
|
for round_id, update in enumerate(result.updates, start=1):
|
|
ledger.record_update(round_id=round_id, model_version="v1", update=update)
|
|
|
|
entries = ledger.list_entries()
|
|
assert len(entries) == 2
|
|
assert len(steps) == len(market.requests)
|
|
assert steps[0].delivery_ns > steps[0].timestamp_ns
|