docs: validate baseline and add deterministic regression coverage

Session 3 baseline validation (no beta/support-weighted code added).

- Two measured full-data original_protean conditions (X-IIoTID and 5G-NIDD,
  alpha=0.5, seed=0); frozen manifests reused and reload-verified.
- Partition consistency, per-class/per-client inspection, round-0, loss audit,
  inference validation, reproducibility and a deterministic end-to-end test.
- S_j participant rule remains an explicitly labelled REPRODUCTION ASSUMPTION;
  |S_j| < M quantified in 53/162 class-condition cells (32.7%), 15/18 conditions.
- Adds docs/baseline/BASELINE_VALIDATION.md, tests/test_baseline_validation.py,
  scripts/validate_sj_participant_rule.py and small analysis artifacts.
- Test suite: 46 passed.

planningawal/ and raw datasets unchanged.

tests/test_baseline_validation.py

@@ -0,0 +1,408 @@
+"""Baseline validation tests (round-0 handling, loss audit, inference, regression).
+
+These tests protect the frozen baseline semantics against accidental change when
+the proposed support-weighted aggregation is added later. They assert robust
+invariants and, where needed, replicate the documented loss formula
+independently rather than calling the production loss inline.
+
+Nothing here implements or exercises the proposed beta/support-weighted method.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from protean_research.federated import TrainingConfig, run_federated  # noqa: E402
+from protean_research.federated import client as client_mod  # noqa: E402
+from protean_research.federated import server as server_mod  # noqa: E402
+from protean_research.federated.client import ClientUpdate, local_update  # noqa: E402
+from protean_research.models.protean_cnn import (  # noqa: E402
+    ModelConfig,
+    average_parameters,
+    build_model,
+    get_parameters,
+    set_parameters,
+)
+from protean_research.models.prototypes import (  # noqa: E402
+    compute_local_prototypes,
+    nearest_prototype_predict,
+)
+
+
+def make_separable(num_classes=3, per_class=12, dim=5, seed=0):
+    rng = np.random.default_rng(seed)
+    features = rng.normal(size=(num_classes * per_class, dim)).astype(np.float32)
+    labels = np.repeat(np.arange(num_classes), per_class).astype(np.int64)
+    features = features + labels[:, None] * 2.0
+    return features, labels
+
+
+# ---------------------------------------------------------------------------
+# Phase 7: round-0 handling and prototype alignment entry
+# ---------------------------------------------------------------------------
+
+
+def test_round0_sees_no_global_prototypes_and_round1_does(monkeypatch):
+    """Round 0 must start with no global prototypes; later rounds receive them."""
+    captured: list = []
+    original = client_mod.local_update
+
+    def spy(*args, **kwargs):
+        # global_prototypes is the 4th positional argument in round.py.
+        captured.append(args[3] if len(args) > 3 else kwargs.get("global_prototypes"))
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(client_mod, "local_update", spy)
+
+    features, labels = make_separable(num_classes=3, per_class=8, dim=8)
+    num_clients = 3
+    assignments = [np.arange(0, 8), np.arange(8, 16), np.arange(16, 24)]
+    config = TrainingConfig(
+        rounds=2,
+        local_epochs=1,
+        batch_size=8,
+        num_clients=num_clients,
+        embedding_dim=8,
+        dense_units=8,
+    )
+    run_federated(
+        train_features=features,
+        train_labels=labels,
+        assignments=assignments,
+        val_features=features,
+        val_labels=labels,
+        num_classes=3,
+        input_dim=features.shape[1],
+        config=config,
+        init_seed_parts=("round0", "seed=0"),
+    )
+
+    assert len(captured) == 2 * num_clients
+    round0 = captured[:num_clients]
+    round1 = captured[num_clients:]
+    # Round 0: the prototype-alignment term is skipped (empty/None prototypes).
+    assert all(g is None or len(g) == 0 for g in round0)
+    # Round 1: the first aggregated global prototypes are broadcast back.
+    assert all(g is not None and len(g) > 0 for g in round1)
+
+
+def test_local_update_without_global_prototypes_still_trains():
+    features, labels = make_separable(num_classes=2, per_class=6, dim=4)
+    model = build_model(ModelConfig(input_dim=4, num_classes=2, embedding_dim=6, dense_units=6))
+    params = get_parameters(model)
+    update = local_update(
+        0,
+        model,
+        params,
+        None,
+        features,
+        labels,
+        np.arange(len(labels)),
+        num_classes=2,
+        local_epochs=1,
+        batch_size=6,
+        learning_rate=1e-3,
+        mu=0.1,
+        lam=1.0,
+        seed=1,
+    )
+    assert isinstance(update, ClientUpdate)
+    assert set(update.prototypes) == {0, 1}
+    assert update.n_samples == len(labels)
+
+
+# ---------------------------------------------------------------------------
+# Phase 8: loss implementation audit (independent replication of the formula)
+# ---------------------------------------------------------------------------
+
+
+def _manual_local_update(
+    model,
+    global_parameters,
+    global_prototypes,
+    features,
+    labels,
+    *,
+    local_epochs,
+    batch_size,
+    learning_rate,
+    mu,
+    lam,
+    seed,
+):
+    """Independent transcription of the documented objective.
+
+    CE(mean) + lam * mean_matched ||phi(x) - Cbar_y||^2 + (mu/2) * sum ||w - w_global||^2
+    """
+    import torch.nn.functional as F
+
+    set_parameters(model, global_parameters)
+    device = next(model.parameters()).device
+    x = torch.as_tensor(features, dtype=torch.float32, device=device)
+    y = torch.as_tensor(np.asarray(labels), dtype=torch.long, device=device)
+    global_tensors = [
+        torch.as_tensor(p, dtype=torch.float32, device=device) for p in global_parameters
+    ]
+    proto_tensors = {
+        int(c): torch.as_tensor(v, dtype=torch.float32, device=device)
+        for c, v in (global_prototypes or {}).items()
+    }
+    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(seed)
+
+    model.train()
+    for _ in range(local_epochs):
+        order = torch.randperm(x.shape[0], generator=generator)
+        for start in range(0, x.shape[0], batch_size):
+            idx = order[start : start + batch_size]
+            xb, yb = x[idx], y[idx]
+            optimizer.zero_grad()
+            log_prob = model(xb)
+            loss = F.nll_loss(log_prob, yb)
+            if proto_tensors:
+                emb = model.embedding(xb)
+                align = torch.zeros((), device=device)
+                count = 0
+                for class_id, proto in proto_tensors.items():
+                    mask = yb == class_id
+                    if mask.any():
+                        diff = emb[mask] - proto.unsqueeze(0)
+                        align = align + (diff ** 2).sum(dim=1).sum()
+                        count += int(mask.sum().item())
+                if count:
+                    loss = loss + lam * align / count
+            if mu:
+                penalty = torch.zeros((), device=device)
+                for param, target in zip(model.parameters(), global_tensors):
+                    penalty = penalty + ((param - target) ** 2).sum()
+                loss = loss + (mu / 2.0) * penalty
+            loss.backward()
+            optimizer.step()
+    return get_parameters(model)
+
+
+def test_local_update_matches_documented_loss_formula():
+    """Same seed/batch order -> identical parameters as the documented formula."""
+    torch.manual_seed(0)
+    dim, num_classes = 8, 3
+    cfg = ModelConfig(input_dim=dim, num_classes=num_classes, embedding_dim=6, dense_units=6)
+    global_model = build_model(cfg)
+    global_parameters = get_parameters(global_model)
+
+    rng = np.random.default_rng(3)
+    features = rng.normal(size=(9, dim)).astype(np.float32)
+    labels = np.array([0, 1, 2, 0, 1, 2, 0, 1, 2], dtype=np.int64)
+    global_prototypes = {
+        0: rng.normal(size=6).astype(np.float32),
+        1: rng.normal(size=6).astype(np.float32),
+    }
+
+    seed = 12345
+    lam, mu = 0.7, 0.1
+
+    impl_model = build_model(cfg)
+    # Dropout draws from the global torch RNG; reseed identically for each path.
+    torch.manual_seed(seed)
+    impl = local_update(
+        0,
+        impl_model,
+        global_parameters,
+        global_prototypes,
+        features,
+        labels,
+        np.arange(len(labels)),
+        num_classes=num_classes,
+        local_epochs=1,
+        batch_size=4,
+        learning_rate=1e-3,
+        mu=mu,
+        lam=lam,
+        seed=seed,
+    )
+
+    ref_model = build_model(cfg)
+    torch.manual_seed(seed)
+    expected = _manual_local_update(
+        ref_model,
+        global_parameters,
+        global_prototypes,
+        features,
+        labels,
+        local_epochs=1,
+        batch_size=4,
+        learning_rate=1e-3,
+        mu=mu,
+        lam=lam,
+        seed=seed,
+    )
+
+    for got, want in zip(impl.parameters, expected):
+        assert torch.allclose(got, want, atol=1e-5, rtol=1e-4), (
+            "local_update parameters diverge from the documented loss formula"
+        )
+
+
+def test_prototype_alignment_skips_classes_without_a_global_prototype():
+    """Labels without a global prototype must not add an alignment term."""
+    torch.manual_seed(1)
+    dim, num_classes = 8, 3
+    cfg = ModelConfig(input_dim=dim, num_classes=num_classes, embedding_dim=5, dense_units=5)
+    global_parameters = get_parameters(build_model(cfg))
+
+    rng = np.random.default_rng(7)
+    features = rng.normal(size=(6, dim)).astype(np.float32)
+    labels = np.array([0, 0, 2, 2, 2, 2], dtype=np.int64)
+    # Only class 0 has a global prototype; class 2 must be skipped entirely.
+    global_prototypes = {0: rng.normal(size=5).astype(np.float32)}
+    seed = 999
+
+    impl_model = build_model(cfg)
+    torch.manual_seed(seed)
+    impl = local_update(
+        0, impl_model, global_parameters, global_prototypes, features, labels,
+        np.arange(len(labels)), num_classes=num_classes, local_epochs=1,
+        batch_size=6, learning_rate=1e-3, mu=0.0, lam=1.0, seed=seed,
+    )
+    ref_model = build_model(cfg)
+    torch.manual_seed(seed)
+    expected = _manual_local_update(
+        ref_model, global_parameters, global_prototypes, features, labels,
+        local_epochs=1, batch_size=6, learning_rate=1e-3, mu=0.0, lam=1.0, seed=seed,
+    )
+    for got, want in zip(impl.parameters, expected):
+        assert torch.allclose(got, want, atol=1e-5, rtol=1e-4)
+
+
+def test_training_config_coefficients_are_frozen():
+    """Guard lambda/mu against drifting away from the frozen protocol."""
+    config = TrainingConfig()
+    assert config.lam == pytest.approx(1.0)
+    assert config.mu == pytest.approx(0.1)
+    assert config.rounds == 10
+    assert config.local_epochs == 3
+    assert config.batch_size == 256
+    assert config.learning_rate == pytest.approx(1e-3)
+
+
+def test_model_aggregation_is_equal_weight_mean():
+    torch.manual_seed(2)
+    cfg = ModelConfig(input_dim=8, num_classes=2, embedding_dim=4, dense_units=4)
+    a = get_parameters(build_model(cfg))
+    b = get_parameters(build_model(cfg))
+    avg = average_parameters([a, b])
+    for pa, pb, pm in zip(a, b, avg):
+        assert torch.allclose(pm, (pa + pb) / 2.0, atol=1e-6)
+    # Equal 1/M averaging: identical lists collapse to themselves.
+    same = average_parameters([a, a, a])
+    for pa, pm in zip(a, same):
+        assert torch.allclose(pm, pa, atol=1e-6)
+
+
+# ---------------------------------------------------------------------------
+# Phase 9: inference validation (synthetic fixture)
+# ---------------------------------------------------------------------------
+
+
+def test_nearest_prototype_predict_maps_noncontiguous_class_ids():
+    prototypes = {2: np.array([0.0, 0.0]), 5: np.array([10.0, 10.0])}
+    embeddings = np.array([[0.5, 0.5], [9.0, 8.5]])
+    predictions = nearest_prototype_predict(embeddings, prototypes)
+    assert predictions.tolist() == [2, 5]
+
+
+def test_nearest_prototype_predict_matches_squared_l2_bruteforce():
+    rng = np.random.default_rng(11)
+    prototypes = {c: rng.normal(size=7) for c in (1, 4, 9)}
+    embeddings = rng.normal(size=(20, 7))
+    expected = []
+    for emb in embeddings:
+        dists = {
+            c: float(((emb - p) ** 2).sum()) for c, p in prototypes.items()
+        }
+        expected.append(min(dists, key=dists.get))
+    got = nearest_prototype_predict(embeddings, prototypes)
+    assert got.tolist() == expected
+
+
+def test_nearest_prototype_predict_raises_on_empty_prototypes():
+    with pytest.raises(ValueError):
+        nearest_prototype_predict(np.zeros((1, 3)), {})
+
+
+def test_predict_global_raises_without_prototypes():
+    class _M:
+        pass
+
+    with pytest.raises(ValueError):
+        server_mod.predict_global(_M(), np.zeros((1, 3)), {})
+
+
+def test_compute_local_prototypes_class_index_mapping():
+    embeddings = np.array([[0.0, 0.0], [2.0, 0.0], [4.0, 4.0]], dtype=np.float32)
+    labels = np.array([1, 1, 3], dtype=np.int64)
+    prototypes, supports = compute_local_prototypes(embeddings, labels, num_classes=5)
+    assert set(prototypes) == {1, 3}
+    assert supports[1] == 2
+    assert prototypes[1] == pytest.approx([1.0, 0.0])
+    assert prototypes[3] == pytest.approx([4.0, 4.0])
+
+
+# ---------------------------------------------------------------------------
+# Phase 10: deterministic end-to-end baseline regression
+# ---------------------------------------------------------------------------
+
+
+def _run_regression(seed_parts):
+    features, labels = make_separable(num_classes=3, per_class=10, dim=6, seed=5)
+    assignments = [
+        np.arange(0, 10),
+        np.arange(10, 20),
+        np.arange(20, 30),
+    ]
+    config = TrainingConfig(
+        rounds=2,
+        local_epochs=1,
+        batch_size=10,
+        num_clients=3,
+        embedding_dim=8,
+        dense_units=8,
+    )
+    return run_federated(
+        train_features=features,
+        train_labels=labels,
+        assignments=assignments,
+        val_features=features,
+        val_labels=labels,
+        num_classes=3,
+        input_dim=features.shape[1],
+        config=config,
+        init_seed_parts=seed_parts,
+    )
+
+
+def test_end_to_end_baseline_is_deterministic_and_learns_separable_fixture():
+    first = _run_regression(("regression", "seed=0"))
+    second = _run_regression(("regression", "seed=0"))
+
+    # Determinism: identical parameters, prototypes and metrics for the same seed.
+    for a, b in zip(first["parameters"], second["parameters"]):
+        assert torch.equal(a, b)
+    assert set(first["global_prototypes"]) == set(second["global_prototypes"]) == {0, 1, 2}
+    for c in first["global_prototypes"]:
+        assert np.array_equal(
+            np.asarray(first["global_prototypes"][c]),
+            np.asarray(second["global_prototypes"][c]),
+        )
+    assert first["history"] == second["history"]
+
+    # Robust invariants (not fragile exact floats): the pipeline runs and the
+    # prototype read-out beats chance on a linearly separable fixture.
+    last = first["history"][-1]
+    assert np.isfinite(last["accuracy"])
+    assert np.isfinite(last["macro_f1"])
+    assert last["accuracy"] > 0.5
+    assert first["snapshot"] is not None