feat: implement reproducible PROTEAN baseline pipeline
This commit is contained in:
willtanoe committed
1 parent
2cc6a9859e
commit
0aa8e2ba79
93 files changed
+8173
-658
No files matched your search
@@ -1,20 +1,18 @@
|
||||
"""Server-side prototype aggregation.
|
||||
|
||||
This package is the single locus of the proposed modification. Local model
|
||||
training, the embedding architecture, the loss terms and the inference rule are
|
||||
intentionally left untouched (see docs/research/IMPLEMENTATION_PLAN.md).
|
||||
This package is the single locus of the proposed modification. In the baseline
|
||||
checkpoint it implements only the reproduced original-PROTEAN aggregation. The
|
||||
proposed support-weighted aggregation is deferred to the beta/controls phase.
|
||||
"""
|
||||
|
||||
from .prototype_aggregation import (
|
||||
AggregationMode,
|
||||
aggregate_class_prototypes,
|
||||
aggregate_prototypes,
|
||||
aggregation_weights,
|
||||
beta_weights,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AggregationMode",
|
||||
"aggregate_class_prototypes",
|
||||
"aggregate_prototypes",
|
||||
"aggregation_weights",
|
||||
"beta_weights",
|
||||
]
|
||||
@@ -1,73 +1,40 @@
|
||||
"""Pure server-side prototype aggregation.
|
||||
"""Baseline server-side prototype aggregation.
|
||||
|
||||
Formulation under study
|
||||
-----------------------
|
||||
For a class ``j`` with contributing-client set ``S_j = {i | N_ij > 0 and C_ij exists}``::
|
||||
This module implements **only the reproduced original-PROTEAN prototype
|
||||
aggregation** used by the baseline. The proposed support-weighted aggregation
|
||||
(``N_ij ** beta``) is deliberately NOT present in this checkpoint; it belongs to
|
||||
the later beta/controls phase and must not be introduced before the baseline is
|
||||
verified.
|
||||
|
||||
C_bar_j(beta) = sum_{i in S_j} N_ij**beta * C_ij
|
||||
--------------------------------
|
||||
sum_{i in S_j} N_ij**beta
|
||||
Participant rule (reproduction assumption, not verified against code)
|
||||
---------------------------------------------------------------------
|
||||
The published Equation (2) / thesis Eq. (5.5) writes ``C_bar_j = (1/M) * sum_i
|
||||
C_ij`` over all ``M`` clients, but never states how a class that is missing on a
|
||||
client (``N_ij = 0``) is handled. The thesis separately states that "prototypes
|
||||
are shared only for locally observed classes", so this baseline assumes the
|
||||
server receives prototypes only from the contributor set::
|
||||
|
||||
with candidate ``beta in {0.0, 0.25, 0.5, 0.75, 1.0}``.
|
||||
S_j = { i | N_ij > 0 and a valid local prototype C_ij exists }
|
||||
|
||||
The functions in this module are pure and dependency-free so they can be unit
|
||||
tested in a bare interpreter. They implement only the two *unambiguous* modes:
|
||||
|
||||
``equal_available``
|
||||
equal weight for every contributing client in ``S_j`` (``beta == 0`` control).
|
||||
|
||||
``support_weighted``
|
||||
``w_ij = N_ij ** beta`` normalized over ``S_j`` (the proposed mechanism).
|
||||
|
||||
``original_protean`` is deliberately NOT implemented here. The published
|
||||
Equation (2) writes ``(1/M) * sum_{i=1}^{M} C_ij`` but the paper does not define
|
||||
what happens to ``C_ij`` when ``N_ij == 0`` (omitted / zero vector / stale value /
|
||||
different denominator). Encoding any choice here would silently assume baseline
|
||||
semantics that must first be verified from official code. See
|
||||
``docs/baseline/PROTEAN_BASELINE_INVESTIGATION.md``.
|
||||
and averages them with **equal weights ``1/|S_j|``**. This is an explicit
|
||||
REPRODUCTION ASSUMPTION of this project, not a source-confirmed fact. It must be
|
||||
validated in the baseline-verification phase; see
|
||||
``docs/baseline/BASELINE_IMPLEMENTATION.md``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from enum import Enum
|
||||
from typing import Optional, Sequence
|
||||
from typing import Mapping, Sequence
|
||||
|
||||
Vector = Sequence[float]
|
||||
|
||||
|
||||
class AggregationMode(str, Enum):
|
||||
"""Supported prototype-aggregation modes."""
|
||||
"""Prototype-aggregation modes implemented in the baseline checkpoint."""
|
||||
|
||||
ORIGINAL_PROTEAN = "original_protean"
|
||||
SUPPORT_WEIGHTED = "support_weighted"
|
||||
EQUAL_AVAILABLE = "equal_available"
|
||||
|
||||
|
||||
def _validate_beta(beta: Optional[float]) -> float:
|
||||
if beta is None:
|
||||
raise ValueError("beta is required")
|
||||
beta = float(beta)
|
||||
if math.isnan(beta) or math.isinf(beta):
|
||||
raise ValueError(f"beta must be finite, got {beta!r}")
|
||||
if not 0.0 <= beta <= 1.0:
|
||||
raise ValueError(f"beta must be in the closed interval [0, 1], got {beta!r}")
|
||||
return beta
|
||||
|
||||
|
||||
def _validate_supports(supports: Optional[Sequence[float]]) -> list[float]:
|
||||
if supports is None:
|
||||
raise ValueError("supports (N_ij) are required")
|
||||
result = [float(n) for n in supports]
|
||||
if not result:
|
||||
raise ValueError("supports must not be empty")
|
||||
for n in result:
|
||||
if math.isnan(n) or math.isinf(n) or n <= 0.0:
|
||||
raise ValueError(
|
||||
"supports must be finite and strictly positive: S_j contains "
|
||||
f"only clients with N_ij > 0, got {n!r}"
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def _validate_vectors(prototypes: Sequence[Vector]) -> int:
|
||||
@@ -89,75 +56,61 @@ def _validate_vectors(prototypes: Sequence[Vector]) -> int:
|
||||
return dim
|
||||
|
||||
|
||||
def beta_weights(supports: Sequence[float], beta: float) -> list[float]:
|
||||
"""Return the unnormalized weights ``N_ij ** beta``."""
|
||||
beta = _validate_beta(beta)
|
||||
return [n ** beta for n in _validate_supports(supports)]
|
||||
|
||||
|
||||
def aggregation_weights(
|
||||
prototypes: Sequence[Vector],
|
||||
mode: AggregationMode | str,
|
||||
beta: Optional[float] = None,
|
||||
supports: Optional[Sequence[float]] = None,
|
||||
) -> list[float]:
|
||||
"""Return normalized aggregation weights that sum to 1.
|
||||
|
||||
``prototypes`` is used for length/dimension validation and to mirror the
|
||||
required shape of the weight vector. Weights are returned in the same client
|
||||
order as ``prototypes``.
|
||||
"""
|
||||
_validate_vectors(prototypes)
|
||||
mode = AggregationMode(mode)
|
||||
|
||||
if mode is AggregationMode.ORIGINAL_PROTEAN:
|
||||
raise NotImplementedError(
|
||||
"original_protean aggregation semantics are not yet verified. "
|
||||
"Missing-class behavior (N_ij == 0) is undefined in the paper and no "
|
||||
"official implementation was found. Verify the baseline first; see "
|
||||
"docs/baseline/PROTEAN_BASELINE_INVESTIGATION.md."
|
||||
)
|
||||
|
||||
if mode is AggregationMode.SUPPORT_WEIGHTED:
|
||||
weights = beta_weights(_validate_supports(supports), _validate_beta(beta))
|
||||
else: # EQUAL_AVAILABLE
|
||||
weights = [1.0] * len(prototypes)
|
||||
|
||||
total = sum(weights)
|
||||
if not math.isfinite(total) or total <= 0.0:
|
||||
raise ValueError(f"aggregation weights must sum to a positive value, got {total!r}")
|
||||
return [w / total for w in weights]
|
||||
|
||||
|
||||
def aggregate_prototypes(
|
||||
prototypes: Sequence[Vector],
|
||||
mode: AggregationMode | str,
|
||||
beta: Optional[float] = None,
|
||||
supports: Optional[Sequence[float]] = None,
|
||||
mode: AggregationMode | str = AggregationMode.ORIGINAL_PROTEAN,
|
||||
) -> list[float]:
|
||||
"""Aggregate ``C_ij`` for one class into ``C_bar_j``.
|
||||
"""Aggregate the local prototypes ``C_ij`` of one class into ``C_bar_j``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
prototypes:
|
||||
Local prototypes ``C_ij`` for the contributing clients in ``S_j``.
|
||||
mode:
|
||||
One of :class:`AggregationMode`.
|
||||
beta:
|
||||
Weighting exponent, required for ``support_weighted`` (and meaningful
|
||||
for ``equal_available`` as a documented ``beta == 0`` equivalent).
|
||||
supports:
|
||||
Sample counts ``N_ij`` (> 0), required for ``support_weighted``.
|
||||
Only :attr:`AggregationMode.ORIGINAL_PROTEAN` is implemented in this
|
||||
baseline checkpoint. Other modes raise ``NotImplementedError``.
|
||||
|
||||
The baseline uses equal weights over the supplied contributor set, so the
|
||||
result is the arithmetic mean of ``prototypes``.
|
||||
"""
|
||||
# ``AggregationMode`` contains only the baseline mode, so any other value
|
||||
# raises ``ValueError`` here. This prevents the proposed support-weighted
|
||||
# aggregation from being used before the baseline is verified.
|
||||
AggregationMode(mode)
|
||||
dim = _validate_vectors(prototypes)
|
||||
if supports is not None and len(supports) != len(prototypes):
|
||||
raise ValueError(
|
||||
f"supports length ({len(supports)}) must match prototypes "
|
||||
f"({len(prototypes)})"
|
||||
)
|
||||
weights = aggregation_weights(prototypes, mode, beta=beta, supports=supports)
|
||||
count = len(prototypes)
|
||||
result = [0.0] * dim
|
||||
for weight, vec in zip(weights, prototypes):
|
||||
for vec in prototypes:
|
||||
for d in range(dim):
|
||||
result[d] += weight * float(vec[d])
|
||||
result[d] += float(vec[d]) / count
|
||||
return result
|
||||
|
||||
|
||||
def aggregate_class_prototypes(
|
||||
contributions,
|
||||
mode: AggregationMode | str = AggregationMode.ORIGINAL_PROTEAN,
|
||||
):
|
||||
"""Aggregate per-class local prototypes into global prototypes.
|
||||
|
||||
``contributions`` maps ``class_id -> sequence of (C_ij, N_ij)``. The
|
||||
contributing set ``S_j`` is constructed here by dropping any entry with a
|
||||
missing prototype (``None``) or ``N_ij <= 0``; a missing class is never
|
||||
fabricated. Classes with no valid contributor are omitted from the result.
|
||||
|
||||
Returns ``{class_id: list[float]}`` with one entry per aggregated class.
|
||||
"""
|
||||
if not isinstance(contributions, Mapping):
|
||||
raise TypeError("contributions must be a mapping class_id -> [(C_ij, N_ij), ...]")
|
||||
result = {}
|
||||
for class_id in sorted(contributions, key=lambda c: int(c)):
|
||||
member_prototypes: list[Vector] = []
|
||||
for proto, support in contributions[class_id]:
|
||||
if proto is None:
|
||||
continue
|
||||
if support is None or float(support) <= 0.0:
|
||||
continue
|
||||
member_prototypes.append(proto)
|
||||
if not member_prototypes:
|
||||
continue
|
||||
result[int(class_id)] = aggregate_prototypes(member_prototypes, mode)
|
||||
return result
|
||||
@@ -0,0 +1 @@
|
||||
"""Data loading, preprocessing, splitting and non-IID partitioning."""
|
||||
@@ -0,0 +1,66 @@
|
||||
"""Non-IID characterization of a partition's class-support matrix ``N_ij``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Sequence
|
||||
|
||||
|
||||
def characterize_support(support, class_names: Sequence[str]) -> list[dict]:
|
||||
"""Per-class inequality statistics of ``N_ij``.
|
||||
|
||||
Returns one row per class with totals, contributor counts and the inequality
|
||||
measures required by the protocol (min/max/mean/std/CV/max-min ratio).
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
support = np.asarray(support)
|
||||
if support.ndim != 2:
|
||||
raise ValueError("support must be a 2D [clients, classes] matrix")
|
||||
if support.shape[1] != len(class_names):
|
||||
raise ValueError("class_names length must match support columns")
|
||||
|
||||
rows: list[dict] = []
|
||||
for c, name in enumerate(class_names):
|
||||
column = support[:, c]
|
||||
positive = column[column > 0]
|
||||
n_clients = int((column > 0).sum())
|
||||
mean = float(positive.mean()) if positive.size else 0.0
|
||||
std = float(positive.std(ddof=0)) if positive.size else 0.0
|
||||
cv = float(std / mean) if mean > 0 else 0.0
|
||||
max_min = (
|
||||
float(positive.max() / positive.min()) if positive.size and positive.min() > 0 else 0.0
|
||||
)
|
||||
rows.append(
|
||||
{
|
||||
"class_id": c,
|
||||
"class_name": name,
|
||||
"total": int(column.sum()),
|
||||
"num_clients": n_clients,
|
||||
"n_missing_clients": int((column == 0).sum()),
|
||||
"min": int(positive.min()) if positive.size else 0,
|
||||
"max": int(positive.max()) if positive.size else 0,
|
||||
"mean": mean,
|
||||
"std": std,
|
||||
"cv": cv,
|
||||
"max_min_ratio": max_min,
|
||||
}
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def client_totals(support, client_ids: Sequence[int] | None = None) -> list[dict]:
|
||||
"""Per-client total sample counts and number of missing classes."""
|
||||
import numpy as np
|
||||
|
||||
support = np.asarray(support)
|
||||
rows: list[dict] = []
|
||||
for client in range(support.shape[0]):
|
||||
row = support[client]
|
||||
rows.append(
|
||||
{
|
||||
"client_id": int(client_ids[client]) if client_ids else client,
|
||||
"n_samples": int(row.sum()),
|
||||
"n_missing_classes": int((row == 0).sum()),
|
||||
}
|
||||
)
|
||||
return rows
|
||||
@@ -0,0 +1,225 @@
|
||||
"""Raw dataset loading and cleaning (read-only with respect to the source CSV).
|
||||
|
||||
The raw datasets are external, immutable source data. This module only reads
|
||||
them. Cleaning here is limited to label canonicalisation, duplicate removal and
|
||||
identification of feature columns; numeric/categorical encoding lives in
|
||||
:mod:`protean_research.data.preprocessing`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Mapping
|
||||
|
||||
from ..utils.config import (
|
||||
expand_path,
|
||||
load_datasets_config,
|
||||
resolve_raw_path,
|
||||
)
|
||||
|
||||
LABEL_ALIAS_KEYS = ("fine_grained", "binary", "tool")
|
||||
|
||||
|
||||
def canonical_label(value: object) -> str:
|
||||
"""Canonicalise a class label.
|
||||
|
||||
Strips surrounding whitespace, collapses internal whitespace runs and
|
||||
removes whitespace around underscores. This fixes the X-IIoTID typo
|
||||
``"Lateral _movement"`` -> ``"Lateral_movement"``.
|
||||
"""
|
||||
text = str(value).strip()
|
||||
text = re.sub(r"\s+", " ", text)
|
||||
text = re.sub(r"\s*_\s*", "_", text)
|
||||
return text
|
||||
|
||||
|
||||
def _string_tuple(values, field: str) -> tuple[str, ...]:
|
||||
result = []
|
||||
for value in values:
|
||||
if not isinstance(value, str):
|
||||
raise TypeError(
|
||||
f"{field} entries must be strings; got {value!r}. "
|
||||
"Quote YAML values that contain ':' (e.g. \"Unnamed: 0\")."
|
||||
)
|
||||
result.append(value)
|
||||
return tuple(result)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DatasetSpec:
|
||||
"""Validated configuration for one dataset."""
|
||||
|
||||
key: str
|
||||
name: str
|
||||
label_column: str
|
||||
label_aliases: tuple[str, ...]
|
||||
benign_label: str
|
||||
drop_labels: tuple[str, ...]
|
||||
missing_tokens: tuple[str, ...]
|
||||
drop_duplicate_rows: bool
|
||||
categorical_columns: tuple[str, ...]
|
||||
non_feature_columns: tuple[str, ...]
|
||||
raw_path: Path
|
||||
env_override: str | None = None
|
||||
aliases: Mapping[str, str] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def non_feature_set(self) -> set[str]:
|
||||
return (
|
||||
set(self.non_feature_columns)
|
||||
| {self.label_column}
|
||||
| set(self.label_aliases)
|
||||
)
|
||||
|
||||
|
||||
def list_dataset_keys(datasets_cfg: Mapping[str, Any] | None = None) -> list[str]:
|
||||
cfg = datasets_cfg or load_datasets_config()
|
||||
return list(cfg.get("datasets", {}).keys())
|
||||
|
||||
|
||||
def load_spec(
|
||||
key: str,
|
||||
datasets_cfg: Mapping[str, Any] | None = None,
|
||||
) -> DatasetSpec:
|
||||
"""Build a :class:`DatasetSpec` for dataset *key*.
|
||||
|
||||
Raises ``KeyError`` with a clear message if the key or a required field is
|
||||
missing, rather than silently defaulting.
|
||||
"""
|
||||
cfg = datasets_cfg or load_datasets_config()
|
||||
datasets = cfg.get("datasets", {})
|
||||
if key not in datasets:
|
||||
raise KeyError(
|
||||
f"unknown dataset {key!r}; available: {sorted(datasets)}"
|
||||
)
|
||||
entry = datasets[key]
|
||||
aliases = entry.get("label_column_aliases", {}) or {}
|
||||
alias_values = tuple(
|
||||
str(aliases[k]) for k in LABEL_ALIAS_KEYS if k in aliases
|
||||
)
|
||||
spec = DatasetSpec(
|
||||
key=key,
|
||||
name=str(entry.get("name", key)),
|
||||
label_column=str(entry["label_column"]),
|
||||
label_aliases=alias_values,
|
||||
benign_label=canonical_label(entry["benign_label"]),
|
||||
drop_labels=tuple(canonical_label(v) for v in entry.get("drop_labels", [])),
|
||||
missing_tokens=tuple(str(t) for t in entry.get("missing_tokens", [])),
|
||||
drop_duplicate_rows=bool(entry.get("drop_duplicate_rows", True)),
|
||||
categorical_columns=_string_tuple(entry.get("categorical_columns", []), "categorical_columns"),
|
||||
non_feature_columns=_string_tuple(entry.get("non_feature_columns", []), "non_feature_columns"),
|
||||
raw_path=resolve_raw_path(entry),
|
||||
env_override=entry.get("env_override"),
|
||||
aliases=aliases,
|
||||
)
|
||||
return spec
|
||||
|
||||
|
||||
def resolve_spec_raw_path(spec: DatasetSpec) -> Path:
|
||||
"""Return the resolved raw path, honouring the env override."""
|
||||
import os
|
||||
|
||||
if spec.env_override and os.environ.get(spec.env_override):
|
||||
return expand_path(os.environ[spec.env_override])
|
||||
return spec.raw_path
|
||||
|
||||
|
||||
@dataclass
|
||||
class CleanedData:
|
||||
"""Cleaned dataset: features, integer labels, class metadata."""
|
||||
|
||||
spec: DatasetSpec
|
||||
features: Any # pandas.DataFrame
|
||||
labels: Any # numpy.ndarray[int64]
|
||||
class_names: list[str]
|
||||
benign_index: int
|
||||
feature_columns: list[str]
|
||||
categorical_columns: list[str]
|
||||
numeric_columns: list[str]
|
||||
n_rows_raw: int
|
||||
n_duplicates_dropped: int
|
||||
n_rows_dropped_label: int
|
||||
|
||||
|
||||
def read_raw(spec: DatasetSpec, nrows: int | None = None):
|
||||
"""Read the raw CSV. Pass ``nrows`` only for smoke tests / inspection."""
|
||||
import pandas as pd
|
||||
|
||||
path = resolve_spec_raw_path(spec)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"raw dataset for {spec.key!r} not found at {path}. "
|
||||
"Set the environment override or fix configs/datasets/datasets.yaml."
|
||||
)
|
||||
frame = pd.read_csv(path, low_memory=False, nrows=nrows)
|
||||
# Normalise the unnamed first column header pandas assigns to 5G-NIDD.
|
||||
frame.columns = [str(c).strip() for c in frame.columns]
|
||||
return frame
|
||||
|
||||
|
||||
def clean_frame(frame, spec: DatasetSpec) -> CleanedData:
|
||||
"""Apply the frozen cleaning recipe and return a :class:`CleanedData`."""
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
if spec.label_column not in frame.columns:
|
||||
raise KeyError(
|
||||
f"label column {spec.label_column!r} not found in {spec.key!r} dataset"
|
||||
)
|
||||
|
||||
n_rows_raw = int(len(frame))
|
||||
duplicates = int(frame.duplicated().sum())
|
||||
if spec.drop_duplicate_rows:
|
||||
frame = frame.drop_duplicates(keep="first").reset_index(drop=True)
|
||||
|
||||
labels = frame[spec.label_column].map(canonical_label)
|
||||
mask = ~labels.isin(set(spec.drop_labels))
|
||||
n_dropped_label = int((~mask).sum())
|
||||
frame = frame.loc[mask].reset_index(drop=True)
|
||||
labels = labels.loc[mask].reset_index(drop=True)
|
||||
|
||||
feature_columns = [
|
||||
c for c in frame.columns if c not in spec.non_feature_set
|
||||
]
|
||||
class_names = sorted(labels.unique().tolist())
|
||||
if spec.benign_label not in class_names:
|
||||
raise ValueError(
|
||||
f"benign label {spec.benign_label!r} absent from {spec.key!r} after cleaning; "
|
||||
f"present classes: {class_names}"
|
||||
)
|
||||
class_to_index = {name: i for i, name in enumerate(class_names)}
|
||||
labels_int = labels.map(class_to_index).to_numpy(dtype=np.int64)
|
||||
|
||||
categorical_columns = [
|
||||
c for c in spec.categorical_columns if c in feature_columns
|
||||
]
|
||||
numeric_columns = [
|
||||
c for c in feature_columns if c not in set(categorical_columns)
|
||||
]
|
||||
|
||||
return CleanedData(
|
||||
spec=spec,
|
||||
features=frame[feature_columns].reset_index(drop=True),
|
||||
labels=labels_int,
|
||||
class_names=class_names,
|
||||
benign_index=class_to_index[spec.benign_label],
|
||||
feature_columns=feature_columns,
|
||||
categorical_columns=categorical_columns,
|
||||
numeric_columns=numeric_columns,
|
||||
n_rows_raw=n_rows_raw,
|
||||
n_duplicates_dropped=duplicates,
|
||||
n_rows_dropped_label=n_dropped_label,
|
||||
)
|
||||
|
||||
|
||||
def load_cleaned(
|
||||
key: str,
|
||||
datasets_cfg: Mapping[str, Any] | None = None,
|
||||
nrows: int | None = None,
|
||||
) -> CleanedData:
|
||||
"""Convenience: resolve spec, read raw and clean in one call."""
|
||||
spec = load_spec(key, datasets_cfg)
|
||||
frame = read_raw(spec, nrows=nrows)
|
||||
return clean_frame(frame, spec)
|
||||
@@ -0,0 +1,159 @@
|
||||
"""Partition manifest persistence and exact reproduction.
|
||||
|
||||
Manifests are small, deterministic, version-controlled JSON records. They store
|
||||
the partition recipe (algorithm id + RNG seed), the split description and the
|
||||
resulting class-support matrix ``N_ij`` plus a content hash. Full sample-index
|
||||
arrays are regenerated deterministically and verified against the stored hash,
|
||||
which keeps manifests small while making reload exact.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import platform
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Mapping
|
||||
|
||||
from ..utils.integrity import hash_index_arrays, hash_json
|
||||
from . import partition as partition_mod
|
||||
from . import splits as splits_mod
|
||||
|
||||
MANIFEST_VERSION = "1"
|
||||
|
||||
|
||||
def manifest_path(partition_dir: str | Path, dataset_key: str, alpha: float, seed: int) -> Path:
|
||||
return Path(partition_dir) / f"{dataset_key}__alpha{alpha:g}__seed{seed}.json"
|
||||
|
||||
|
||||
def build_manifest(
|
||||
*,
|
||||
dataset_key: str,
|
||||
dataset_name: str,
|
||||
alpha: float,
|
||||
seed: int,
|
||||
num_clients: int,
|
||||
class_names: list[str],
|
||||
benign_class: str,
|
||||
labels,
|
||||
split: splits_mod.SplitIndices,
|
||||
assignments,
|
||||
rng_seed: int,
|
||||
train_fraction: float,
|
||||
validation_fraction: float,
|
||||
test_fraction: float,
|
||||
config_hash: str | None = None,
|
||||
preprocessing_hash: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
import numpy as np
|
||||
|
||||
labels = np.asarray(labels)
|
||||
num_classes = len(class_names)
|
||||
support = partition_mod.support_matrix(
|
||||
assignments, labels[split.train], num_classes
|
||||
)
|
||||
manifest = {
|
||||
"manifest_version": MANIFEST_VERSION,
|
||||
"dataset": dataset_key,
|
||||
"dataset_name": dataset_name,
|
||||
"alpha": float(alpha),
|
||||
"seed": int(seed),
|
||||
"num_clients": int(num_clients),
|
||||
"num_classes": num_classes,
|
||||
"class_names": list(class_names),
|
||||
"benign_class": benign_class,
|
||||
"algorithm": partition_mod.ALGORITHM,
|
||||
"rng_seed": int(rng_seed),
|
||||
"num_train": int(len(split.train)),
|
||||
"support": support.astype(int).tolist(),
|
||||
"partition_hash": partition_mod.partition_hash(assignments),
|
||||
"split": {
|
||||
"strategy": "stratified_by_class",
|
||||
"seed": int(seed),
|
||||
"train_fraction": float(train_fraction),
|
||||
"validation_fraction": float(validation_fraction),
|
||||
"test_fraction": float(test_fraction),
|
||||
"sizes": split.sizes,
|
||||
"hash": hash_index_arrays([split.train, split.validation, split.test]),
|
||||
},
|
||||
"preprocessing_hash": preprocessing_hash,
|
||||
"config_hash": config_hash,
|
||||
"created_utc": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
|
||||
"python": platform.python_version(),
|
||||
}
|
||||
return manifest
|
||||
|
||||
|
||||
def save_manifest(path: str | Path, manifest: Mapping[str, Any]) -> Path:
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with path.open("w", encoding="utf-8") as handle:
|
||||
json.dump(manifest, handle, indent=2, sort_keys=True)
|
||||
handle.write("\n")
|
||||
return path
|
||||
|
||||
|
||||
def load_manifest(path: str | Path) -> dict[str, Any]:
|
||||
with Path(path).open("r", encoding="utf-8") as handle:
|
||||
manifest = json.load(handle)
|
||||
if manifest.get("manifest_version") != MANIFEST_VERSION:
|
||||
raise ValueError(
|
||||
f"unsupported manifest version {manifest.get('manifest_version')!r}"
|
||||
)
|
||||
return manifest
|
||||
|
||||
|
||||
def reproduce_split(manifest: Mapping[str, Any], labels):
|
||||
"""Recompute the split described by *manifest* and verify its hash."""
|
||||
split_cfg = manifest["split"]
|
||||
split = splits_mod.stratified_split(
|
||||
labels,
|
||||
dataset_key=str(manifest["dataset"]),
|
||||
seed=int(split_cfg["seed"]),
|
||||
train_fraction=float(split_cfg["train_fraction"]),
|
||||
validation_fraction=float(split_cfg["validation_fraction"]),
|
||||
test_fraction=float(split_cfg["test_fraction"]),
|
||||
)
|
||||
expected_hash = split_cfg["hash"]
|
||||
actual_hash = hash_index_arrays([split.train, split.validation, split.test])
|
||||
if actual_hash != expected_hash:
|
||||
raise ValueError(
|
||||
"split mismatch: manifest does not reproduce on this dataset; "
|
||||
"the raw data or split recipe changed"
|
||||
)
|
||||
if split.sizes != split_cfg["sizes"]:
|
||||
raise ValueError(
|
||||
f"split sizes mismatch: {split.sizes} != {split_cfg['sizes']}"
|
||||
)
|
||||
return split
|
||||
|
||||
|
||||
def reproduce_train_assignments(manifest: Mapping[str, Any], labels):
|
||||
"""Regenerate train assignments and verify ``N_ij`` and content hash."""
|
||||
import numpy as np
|
||||
|
||||
split = reproduce_split(manifest, labels)
|
||||
labels_train = np.asarray(labels)[split.train]
|
||||
class_indices = partition_mod.class_index_map(labels_train, int(manifest["num_classes"]))
|
||||
assignments = partition_mod.dirichlet_partition(
|
||||
class_indices,
|
||||
int(manifest["num_clients"]),
|
||||
float(manifest["alpha"]),
|
||||
int(manifest["rng_seed"]),
|
||||
)
|
||||
support = partition_mod.support_matrix(
|
||||
assignments, labels_train, int(manifest["num_classes"])
|
||||
)
|
||||
expected_support = np.asarray(manifest["support"], dtype=np.int64)
|
||||
if not np.array_equal(support, expected_support):
|
||||
raise ValueError("support matrix N_ij mismatch on manifest reload")
|
||||
if partition_mod.partition_hash(assignments) != manifest["partition_hash"]:
|
||||
raise ValueError("partition content hash mismatch on manifest reload")
|
||||
return split, assignments, support
|
||||
|
||||
|
||||
def manifest_content_hash(manifest: Mapping[str, Any]) -> str:
|
||||
payload = dict(manifest)
|
||||
payload.pop("created_utc", None)
|
||||
payload.pop("python", None)
|
||||
return hash_json(payload)
|
||||
@@ -0,0 +1,129 @@
|
||||
"""Dirichlet non-IID client partitioning.
|
||||
|
||||
Frozen algorithm (``docs/research/EXPERIMENT_PROTOCOL.md`` §2): for each class,
|
||||
draw client proportions from ``Dirichlet(alpha)`` and multinomial-allocate that
|
||||
class's samples. Partitions are generated from the training split only and are
|
||||
persisted in a manifest; a beta run must load the manifest, never regenerate.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from ..utils.integrity import hash_index_arrays
|
||||
from ..utils.seeding import derive_seed
|
||||
|
||||
ALGORITHM = "per_class_dirichlet_multinomial"
|
||||
|
||||
|
||||
def partition_rng_seed(dataset_key: str, alpha: float, seed: int, tag: str = "dirichlet") -> int:
|
||||
"""Deterministic RNG seed for a ``(dataset, alpha, seed)`` partition."""
|
||||
return derive_seed(dataset_key, tag, f"alpha={alpha}", f"seed={seed}")
|
||||
|
||||
|
||||
def class_index_map(labels, num_classes: int) -> dict[int, object]:
|
||||
"""Map class id -> positional indices within *labels*."""
|
||||
import numpy as np
|
||||
|
||||
labels = np.asarray(labels)
|
||||
return {
|
||||
c: np.flatnonzero(labels == c).astype(np.int64)
|
||||
for c in range(num_classes)
|
||||
}
|
||||
|
||||
|
||||
def dirichlet_partition(
|
||||
class_indices: dict[int, object],
|
||||
num_clients: int,
|
||||
alpha: float,
|
||||
rng_seed: int,
|
||||
) -> list[object]:
|
||||
"""Allocate samples to clients and return a client->indices list.
|
||||
|
||||
``class_indices`` maps class id -> positions. The output list has length
|
||||
``num_clients``; each element is an ``int64`` array (possibly empty).
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
if num_clients <= 0:
|
||||
raise ValueError("num_clients must be positive")
|
||||
if alpha <= 0:
|
||||
raise ValueError("alpha must be positive")
|
||||
|
||||
rng = np.random.default_rng(rng_seed)
|
||||
buckets: list[list[np.ndarray]] = [[] for _ in range(num_clients)]
|
||||
|
||||
for class_id in sorted(class_indices):
|
||||
positions = np.asarray(class_indices[class_id], dtype=np.int64)
|
||||
n = len(positions)
|
||||
if n == 0:
|
||||
continue
|
||||
shuffled = positions[rng.permutation(n)]
|
||||
proportions = rng.dirichlet(np.full(num_clients, alpha, dtype=float))
|
||||
counts = rng.multinomial(n, proportions)
|
||||
start = 0
|
||||
for client in range(num_clients):
|
||||
k = int(counts[client])
|
||||
if k > 0:
|
||||
buckets[client].append(shuffled[start : start + k])
|
||||
start += k
|
||||
|
||||
return [
|
||||
(
|
||||
np.sort(np.concatenate(parts)).astype(np.int64, copy=False)
|
||||
if parts
|
||||
else np.empty(0, dtype=np.int64)
|
||||
)
|
||||
for parts in buckets
|
||||
]
|
||||
|
||||
|
||||
def support_matrix(assignments, labels, num_classes: int):
|
||||
"""Return the ``[num_clients, num_classes]`` count matrix ``N_ij``."""
|
||||
import numpy as np
|
||||
|
||||
labels = np.asarray(labels)
|
||||
matrix = np.zeros((len(assignments), num_classes), dtype=np.int64)
|
||||
for client, idx in enumerate(assignments):
|
||||
if len(idx):
|
||||
counts = np.bincount(labels[idx], minlength=num_classes)
|
||||
matrix[client, : len(counts)] = counts
|
||||
return matrix
|
||||
|
||||
|
||||
def partition_hash(assignments) -> str:
|
||||
return hash_index_arrays(assignments)
|
||||
|
||||
|
||||
def partition_eval_split(
|
||||
labels,
|
||||
num_clients: int,
|
||||
alpha: float,
|
||||
dataset_key: str,
|
||||
seed: int,
|
||||
split_name: str,
|
||||
):
|
||||
"""Partition a validation/test split into per-client local eval sets.
|
||||
|
||||
The baseline evaluation is per client (each client scores its local held-out
|
||||
data). This uses the same frozen Dirichlet algorithm with a split-specific
|
||||
seed, so it is deterministic and paired across variants.
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
labels = np.asarray(labels)
|
||||
# Labels are globally encoded; class space size is max(label) + 1.
|
||||
num_classes = int(labels.max()) + 1 if len(labels) else 0
|
||||
class_indices = class_index_map(labels, num_classes)
|
||||
rng_seed = derive_seed(dataset_key, f"eval_partition:{split_name}", f"alpha={alpha}", f"seed={seed}")
|
||||
return dirichlet_partition(class_indices, num_clients, alpha, rng_seed), rng_seed
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PartitionView:
|
||||
"""Convenience bundle for a generated partition."""
|
||||
|
||||
assignments: list
|
||||
support: object # np.ndarray [clients, classes]
|
||||
rng_seed: int
|
||||
content_hash: str
|
||||
@@ -0,0 +1,131 @@
|
||||
"""Deterministic data pipeline: raw CSV -> cleaned -> split -> preprocess -> partition.
|
||||
|
||||
This module wires the frozen protocol together so that experiment runners do not
|
||||
duplicate orchestration logic. It never modifies the raw datasets.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Mapping
|
||||
|
||||
from ..utils.integrity import hash_json
|
||||
from . import loading, manifest as manifest_mod, partition as partition_mod, splits as splits_mod
|
||||
from .preprocessing import Preprocessor, fit_preprocessor, transform
|
||||
|
||||
TRAIN_FRACTION = 0.70
|
||||
VALIDATION_FRACTION = 0.10
|
||||
TEST_FRACTION = 0.20
|
||||
|
||||
|
||||
@dataclass
|
||||
class DatasetBundle:
|
||||
"""Fully prepared dataset for a fixed ``(dataset, seed)``."""
|
||||
|
||||
spec: loading.DatasetSpec
|
||||
cleaned: loading.CleanedData
|
||||
split: splits_mod.SplitIndices
|
||||
preprocessor: Preprocessor
|
||||
features: Any # float32 matrix for all cleaned rows
|
||||
train_fraction: float = TRAIN_FRACTION
|
||||
validation_fraction: float = VALIDATION_FRACTION
|
||||
test_fraction: float = TEST_FRACTION
|
||||
|
||||
@property
|
||||
def class_names(self) -> list[str]:
|
||||
return list(self.cleaned.class_names)
|
||||
|
||||
@property
|
||||
def num_classes(self) -> int:
|
||||
return len(self.cleaned.class_names)
|
||||
|
||||
@property
|
||||
def labels(self):
|
||||
return self.cleaned.labels
|
||||
|
||||
|
||||
def prepare_dataset(
|
||||
dataset_key: str,
|
||||
*,
|
||||
seed: int,
|
||||
datasets_cfg: Mapping[str, Any] | None = None,
|
||||
nrows: int | None = None,
|
||||
raw_frame: Any | None = None,
|
||||
) -> DatasetBundle:
|
||||
"""Run the deterministic, leak-free preprocessing pipeline.
|
||||
|
||||
``nrows`` is a smoke-test-only override that limits how many raw rows are
|
||||
read; ``raw_frame`` lets a smoke caller pass an already-subsampled frame.
|
||||
Neither must be used for a scientific run.
|
||||
"""
|
||||
spec = loading.load_spec(dataset_key, datasets_cfg)
|
||||
frame = raw_frame if raw_frame is not None else loading.read_raw(spec, nrows=nrows)
|
||||
cleaned = loading.clean_frame(frame, spec)
|
||||
split = splits_mod.stratified_split(
|
||||
cleaned.labels,
|
||||
dataset_key=dataset_key,
|
||||
seed=seed,
|
||||
train_fraction=TRAIN_FRACTION,
|
||||
validation_fraction=VALIDATION_FRACTION,
|
||||
test_fraction=TEST_FRACTION,
|
||||
)
|
||||
splits_mod.check_disjoint(split, n_rows=len(cleaned.labels))
|
||||
preprocessor = fit_preprocessor(cleaned, split.train)
|
||||
features = transform(preprocessor, cleaned)
|
||||
return DatasetBundle(
|
||||
spec=spec,
|
||||
cleaned=cleaned,
|
||||
split=split,
|
||||
preprocessor=preprocessor,
|
||||
features=features,
|
||||
)
|
||||
|
||||
|
||||
def create_partition(
|
||||
bundle: DatasetBundle,
|
||||
*,
|
||||
alpha: float,
|
||||
seed: int,
|
||||
num_clients: int,
|
||||
partition_dir: str | Path,
|
||||
preprocessing_hash: str | None = None,
|
||||
config_hash: str | None = None,
|
||||
save: bool = True,
|
||||
) -> tuple[dict, list, Any]:
|
||||
"""Create and (optionally) persist the partition manifest for one condition."""
|
||||
import numpy as np
|
||||
|
||||
labels_train = np.asarray(bundle.labels)[bundle.split.train]
|
||||
class_indices = partition_mod.class_index_map(labels_train, bundle.num_classes)
|
||||
rng_seed = partition_mod.partition_rng_seed(bundle.spec.key, alpha, seed)
|
||||
assignments = partition_mod.dirichlet_partition(
|
||||
class_indices, num_clients, alpha, rng_seed
|
||||
)
|
||||
manifest = manifest_mod.build_manifest(
|
||||
dataset_key=bundle.spec.key,
|
||||
dataset_name=bundle.spec.name,
|
||||
alpha=alpha,
|
||||
seed=seed,
|
||||
num_clients=num_clients,
|
||||
class_names=bundle.class_names,
|
||||
benign_class=bundle.spec.benign_label,
|
||||
labels=bundle.labels,
|
||||
split=bundle.split,
|
||||
assignments=assignments,
|
||||
rng_seed=rng_seed,
|
||||
train_fraction=bundle.train_fraction,
|
||||
validation_fraction=bundle.validation_fraction,
|
||||
test_fraction=bundle.test_fraction,
|
||||
preprocessing_hash=preprocessing_hash,
|
||||
config_hash=config_hash,
|
||||
)
|
||||
if save:
|
||||
path = manifest_mod.manifest_path(partition_dir, bundle.spec.key, alpha, seed)
|
||||
manifest_mod.save_manifest(path, manifest)
|
||||
support = np.asarray(manifest["support"], dtype=np.int64)
|
||||
return manifest, assignments, support
|
||||
|
||||
|
||||
def preprocessing_content_hash(preprocessor: Preprocessor) -> str:
|
||||
return hash_json(preprocessor.to_dict())
|
||||
@@ -0,0 +1,187 @@
|
||||
"""Deterministic, leak-free feature preprocessing.
|
||||
|
||||
The frozen recipe is described in ``docs/research/EXPERIMENT_PROTOCOL.md`` §13.
|
||||
All preprocessing statistics (medians, means, stds, one-hot categories) are
|
||||
fitted on the **train split only** and then applied unchanged to validation and
|
||||
test. The fitted state is serialisable so future rounds and beta variants reuse
|
||||
exactly the same transform.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Sequence
|
||||
|
||||
MISSING_CATEGORY = "__MISSING__"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Preprocessor:
|
||||
numeric_columns: list[str]
|
||||
categorical_columns: list[str]
|
||||
medians: dict[str, float]
|
||||
means: dict[str, float]
|
||||
stds: dict[str, float]
|
||||
categories: dict[str, list[str]]
|
||||
output_columns: list[str] = field(default_factory=list)
|
||||
|
||||
# ---- persistence -------------------------------------------------
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"numeric_columns": self.numeric_columns,
|
||||
"categorical_columns": self.categorical_columns,
|
||||
"medians": self.medians,
|
||||
"means": self.means,
|
||||
"stds": self.stds,
|
||||
"categories": self.categories,
|
||||
"output_columns": self.output_columns,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, payload: dict[str, Any]) -> "Preprocessor":
|
||||
return cls(
|
||||
numeric_columns=list(payload["numeric_columns"]),
|
||||
categorical_columns=list(payload["categorical_columns"]),
|
||||
medians=dict(payload["medians"]),
|
||||
means=dict(payload["means"]),
|
||||
stds=dict(payload["stds"]),
|
||||
categories={k: list(v) for k, v in payload["categories"].items()},
|
||||
output_columns=list(payload.get("output_columns", [])),
|
||||
)
|
||||
|
||||
def save(self, path: str | Path) -> None:
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with path.open("w", encoding="utf-8") as handle:
|
||||
json.dump(self.to_dict(), handle, indent=2, sort_keys=True)
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: str | Path) -> "Preprocessor":
|
||||
with Path(path).open("r", encoding="utf-8") as handle:
|
||||
return cls.from_dict(json.load(handle))
|
||||
|
||||
|
||||
def _replace_missing(series, missing_tokens: Sequence[str]):
|
||||
"""Map configured missing tokens to NaN."""
|
||||
import numpy as np
|
||||
|
||||
if not missing_tokens:
|
||||
return series
|
||||
result = series.astype(object)
|
||||
for token in missing_tokens:
|
||||
result = result.mask(result == token, np.nan)
|
||||
return result
|
||||
|
||||
|
||||
def _output_columns(pre: Preprocessor) -> list[str]:
|
||||
cols = list(pre.numeric_columns)
|
||||
for col in pre.categorical_columns:
|
||||
for cat in pre.categories[col]:
|
||||
cols.append(f"{col}=={cat}")
|
||||
return cols
|
||||
|
||||
|
||||
def fit_preprocessor(cleaned, train_idx) -> Preprocessor:
|
||||
"""Fit preprocessing state on the train split only.
|
||||
|
||||
``cleaned`` is a :class:`~protean_research.data.loading.CleanedData`;
|
||||
``train_idx`` are positional indices into ``cleaned.features``.
|
||||
"""
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
features = cleaned.features
|
||||
spec = cleaned.spec
|
||||
train = features.iloc[train_idx]
|
||||
|
||||
medians: dict[str, float] = {}
|
||||
means: dict[str, float] = {}
|
||||
stds: dict[str, float] = {}
|
||||
|
||||
for col in cleaned.numeric_columns:
|
||||
values = pd.to_numeric(
|
||||
_replace_missing(train[col], spec.missing_tokens), errors="coerce"
|
||||
).astype("float64")
|
||||
median = float(np.nanmedian(values.to_numpy())) if values.notna().any() else 0.0
|
||||
median = 0.0 if not np.isfinite(median) else median
|
||||
filled = values.fillna(median).to_numpy(dtype="float64")
|
||||
mean = float(filled.mean()) if filled.size else 0.0
|
||||
std = float(filled.std(ddof=0)) if filled.size else 0.0
|
||||
if not np.isfinite(std) or std == 0.0:
|
||||
std = 1.0
|
||||
medians[col] = median
|
||||
means[col] = mean
|
||||
stds[col] = std
|
||||
|
||||
categories: dict[str, list[str]] = {}
|
||||
for col in cleaned.categorical_columns:
|
||||
values = _replace_missing(train[col], spec.missing_tokens)
|
||||
observed = {str(v) for v in values.dropna().unique().tolist()}
|
||||
cats = sorted(observed)
|
||||
if values.isna().any():
|
||||
cats.append(MISSING_CATEGORY)
|
||||
categories[col] = cats
|
||||
|
||||
pre = Preprocessor(
|
||||
numeric_columns=list(cleaned.numeric_columns),
|
||||
categorical_columns=list(cleaned.categorical_columns),
|
||||
medians=medians,
|
||||
means=means,
|
||||
stds=stds,
|
||||
categories=categories,
|
||||
)
|
||||
pre.output_columns = _output_columns(pre)
|
||||
return pre
|
||||
|
||||
|
||||
def transform(pre: Preprocessor, cleaned, indices=None):
|
||||
"""Apply fitted preprocessing and return a ``float32`` matrix.
|
||||
|
||||
``indices`` selects rows of ``cleaned.features`` (all rows if ``None``).
|
||||
"""
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
features = cleaned.features
|
||||
spec = cleaned.spec
|
||||
frame = features if indices is None else features.iloc[indices]
|
||||
n = len(frame)
|
||||
|
||||
blocks: list[np.ndarray] = []
|
||||
for col in pre.numeric_columns:
|
||||
values = pd.to_numeric(
|
||||
_replace_missing(frame[col], spec.missing_tokens), errors="coerce"
|
||||
).astype("float64")
|
||||
values = values.fillna(pre.medians[col]).to_numpy(dtype="float64")
|
||||
values = (values - pre.means[col]) / pre.stds[col]
|
||||
blocks.append(values.astype(np.float32).reshape(n, 1))
|
||||
|
||||
for col in pre.categorical_columns:
|
||||
raw = _replace_missing(frame[col], spec.missing_tokens)
|
||||
missing = raw.isna().to_numpy()
|
||||
text = raw.astype("object").where(~missing, other=None)
|
||||
cats = pre.categories[col]
|
||||
index = {cat: i for i, cat in enumerate(cats)}
|
||||
onehot = np.zeros((n, len(cats)), dtype=np.float32)
|
||||
for row, value in enumerate(text.to_numpy()):
|
||||
if value is None:
|
||||
if MISSING_CATEGORY in index:
|
||||
onehot[row, index[MISSING_CATEGORY]] = 1.0
|
||||
continue
|
||||
j = index.get(str(value))
|
||||
if j is not None:
|
||||
onehot[row, j] = 1.0
|
||||
blocks.append(onehot)
|
||||
|
||||
if not blocks:
|
||||
raise ValueError("no feature columns available after preprocessing")
|
||||
matrix = np.concatenate(blocks, axis=1)
|
||||
return np.ascontiguousarray(matrix, dtype=np.float32)
|
||||
|
||||
|
||||
def fit_transform(cleaned, train_idx):
|
||||
"""Fit on ``train_idx`` and return ``(preprocessor, full_matrix)``."""
|
||||
pre = fit_preprocessor(cleaned, train_idx)
|
||||
return pre, transform(pre, cleaned)
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Rare-class and unseen-class definitions.
|
||||
|
||||
Frozen research definitions (``planningawal/PROTEAN_Revisi_Lurus.md`` §9):
|
||||
per client, rare classes are the two attack classes with the smallest positive
|
||||
training support; unseen classes are attack classes with ``N_ij = 0``. The
|
||||
benign class is never rare or unseen. Rare and unseen are never merged.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Sequence
|
||||
|
||||
|
||||
def attack_class_ids(num_classes: int, benign_index: int) -> list[int]:
|
||||
return [c for c in range(num_classes) if c != benign_index]
|
||||
|
||||
|
||||
def rare_class_ids_for_client(
|
||||
support_row: Sequence[int],
|
||||
benign_index: int,
|
||||
max_rare: int = 2,
|
||||
) -> list[int]:
|
||||
"""Two smallest positive-support attack classes (ties broken by class id)."""
|
||||
support_row = list(support_row)
|
||||
candidates = [
|
||||
(int(support_row[c]), c)
|
||||
for c in range(len(support_row))
|
||||
if c != benign_index and int(support_row[c]) > 0
|
||||
]
|
||||
candidates.sort(key=lambda item: (item[0], item[1]))
|
||||
return [c for _, c in candidates[:max_rare]]
|
||||
|
||||
|
||||
def unseen_class_ids_for_client(
|
||||
support_row: Sequence[int],
|
||||
benign_index: int,
|
||||
) -> list[int]:
|
||||
support_row = list(support_row)
|
||||
return [
|
||||
c
|
||||
for c in range(len(support_row))
|
||||
if c != benign_index and int(support_row[c]) == 0
|
||||
]
|
||||
|
||||
|
||||
def rare_unseen_rows(
|
||||
support,
|
||||
class_names: Sequence[str],
|
||||
benign_index: int,
|
||||
dataset_key: str,
|
||||
alpha: float,
|
||||
seed: int,
|
||||
) -> list[dict]:
|
||||
"""One row per ``(client, class)`` with rare/unseen flags and support."""
|
||||
import numpy as np
|
||||
|
||||
support = np.asarray(support)
|
||||
num_classes = support.shape[1]
|
||||
rows: list[dict] = []
|
||||
for client in range(support.shape[0]):
|
||||
row = support[client]
|
||||
rare = set(rare_class_ids_for_client(row, benign_index))
|
||||
unseen = set(unseen_class_ids_for_client(row, benign_index))
|
||||
for c in range(num_classes):
|
||||
rows.append(
|
||||
{
|
||||
"dataset": dataset_key,
|
||||
"alpha": float(alpha),
|
||||
"seed": int(seed),
|
||||
"client_id": client,
|
||||
"class_id": c,
|
||||
"class_name": class_names[c],
|
||||
"N_ij": int(row[c]),
|
||||
"rare_flag": c in rare,
|
||||
"unseen_flag": c in unseen,
|
||||
}
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def rare_class_ids(
|
||||
support,
|
||||
benign_index: int,
|
||||
) -> list[list[int]]:
|
||||
"""Per-client list of rare class ids."""
|
||||
import numpy as np
|
||||
|
||||
support = np.asarray(support)
|
||||
return [
|
||||
rare_class_ids_for_client(support[client], benign_index)
|
||||
for client in range(support.shape[0])
|
||||
]
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Deterministic stratified train/validation/test splitting.
|
||||
|
||||
Frozen protocol (``docs/research/EXPERIMENT_PROTOCOL.md`` §3): a single
|
||||
stratified 70/10/20 split per ``(dataset, seed)``, reused by every aggregation
|
||||
variant and beta candidate. Test data never influences preprocessing, model
|
||||
selection or beta selection.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from ..utils.seeding import numpy_rng
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SplitIndices:
|
||||
"""Positional indices into the cleaned dataset for each split."""
|
||||
|
||||
train: object # np.ndarray[int64]
|
||||
validation: object
|
||||
test: object
|
||||
|
||||
@property
|
||||
def sizes(self) -> dict[str, int]:
|
||||
return {
|
||||
"train": int(len(self.train)),
|
||||
"validation": int(len(self.validation)),
|
||||
"test": int(len(self.test)),
|
||||
}
|
||||
|
||||
|
||||
def stratified_split(
|
||||
labels,
|
||||
dataset_key: str,
|
||||
seed: int,
|
||||
train_fraction: float = 0.70,
|
||||
validation_fraction: float = 0.10,
|
||||
test_fraction: float = 0.20,
|
||||
) -> SplitIndices:
|
||||
"""Return a deterministic per-class stratified split.
|
||||
|
||||
Every class is shuffled with the seeded RNG and split into train/validation/
|
||||
test so that all classes are represented in validation and test, which the
|
||||
rare-class evaluation requires.
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
fractions = np.array([train_fraction, validation_fraction, test_fraction], dtype=float)
|
||||
if fractions.min() < 0 or abs(fractions.sum() - 1.0) > 1e-9:
|
||||
raise ValueError(
|
||||
"train/validation/test fractions must be non-negative and sum to 1, "
|
||||
f"got {fractions.tolist()}"
|
||||
)
|
||||
|
||||
labels = np.asarray(labels)
|
||||
rng = numpy_rng(dataset_key, "split", f"seed={seed}")
|
||||
train_parts: list[np.ndarray] = []
|
||||
val_parts: list[np.ndarray] = []
|
||||
test_parts: list[np.ndarray] = []
|
||||
|
||||
for class_id in sorted(np.unique(labels).tolist()):
|
||||
positions = np.flatnonzero(labels == class_id)
|
||||
positions = positions[rng.permutation(len(positions))]
|
||||
n = len(positions)
|
||||
n_val = int(round(n * validation_fraction))
|
||||
n_test = int(round(n * test_fraction))
|
||||
# Keep at least one sample in train when a class is tiny.
|
||||
if n_val + n_test >= n and n >= 1:
|
||||
n_val = min(n_val, n - 1)
|
||||
n_test = min(n_test, n - n_val - 1)
|
||||
val_parts.append(positions[:n_val])
|
||||
test_parts.append(positions[n_val : n_val + n_test])
|
||||
train_parts.append(positions[n_val + n_test :])
|
||||
|
||||
def _concat(parts):
|
||||
if not parts:
|
||||
return np.empty(0, dtype=np.int64)
|
||||
joined = np.concatenate(parts)
|
||||
joined.sort()
|
||||
return joined.astype(np.int64, copy=False)
|
||||
|
||||
return SplitIndices(
|
||||
train=_concat(train_parts),
|
||||
validation=_concat(val_parts),
|
||||
test=_concat(test_parts),
|
||||
)
|
||||
|
||||
|
||||
def check_disjoint(splits: SplitIndices, n_rows: int | None = None) -> None:
|
||||
"""Raise ``ValueError`` if the three splits overlap or miss rows.
|
||||
|
||||
This is a leakage guard used by tests and the smoke runner.
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
train = np.asarray(splits.train)
|
||||
val = np.asarray(splits.validation)
|
||||
test = np.asarray(splits.test)
|
||||
if len(np.intersect1d(train, val)) or len(np.intersect1d(train, test)) or len(
|
||||
np.intersect1d(val, test)
|
||||
):
|
||||
raise ValueError("train/validation/test splits overlap")
|
||||
if n_rows is not None:
|
||||
union = np.union1d(np.union1d(train, val), test)
|
||||
if len(union) != n_rows:
|
||||
raise ValueError(
|
||||
f"splits cover {len(union)} rows but dataset has {n_rows}"
|
||||
)
|
||||
@@ -0,0 +1,17 @@
|
||||
"""Evaluation metrics for the baseline and future beta comparisons."""
|
||||
|
||||
from .metrics import (
|
||||
classification_metrics,
|
||||
confusion_matrix,
|
||||
mean_metric,
|
||||
rare_class_metrics,
|
||||
unseen_class_report,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"classification_metrics",
|
||||
"confusion_matrix",
|
||||
"rare_class_metrics",
|
||||
"unseen_class_report",
|
||||
"mean_metric",
|
||||
]
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Classification metrics (pure NumPy, no scikit-learn dependency)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Mapping, Sequence
|
||||
|
||||
|
||||
def confusion_matrix(y_true, y_pred, num_classes: int):
|
||||
import numpy as np
|
||||
|
||||
y_true = np.asarray(y_true, dtype=np.int64)
|
||||
y_pred = np.asarray(y_pred, dtype=np.int64)
|
||||
if y_true.shape != y_pred.shape:
|
||||
raise ValueError("y_true and y_pred must have the same shape")
|
||||
matrix = np.zeros((num_classes, num_classes), dtype=np.int64)
|
||||
valid = (y_true >= 0) & (y_true < num_classes) & (y_pred >= 0) & (y_pred < num_classes)
|
||||
np.add.at(matrix, (y_true[valid], y_pred[valid]), 1)
|
||||
return matrix
|
||||
|
||||
|
||||
def _prf_from_confusion(matrix, class_id: int) -> tuple[float, float, float, int]:
|
||||
import numpy as np
|
||||
|
||||
tp = float(matrix[class_id, class_id])
|
||||
fp = float(matrix[:, class_id].sum() - tp)
|
||||
fn = float(matrix[class_id, :].sum() - tp)
|
||||
support = int(matrix[class_id, :].sum())
|
||||
precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
|
||||
recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
|
||||
f1 = (
|
||||
2 * precision * recall / (precision + recall)
|
||||
if (precision + recall) > 0
|
||||
else 0.0
|
||||
)
|
||||
return precision, recall, f1, support
|
||||
|
||||
|
||||
def classification_metrics(y_true, y_pred, num_classes: int) -> dict:
|
||||
"""Accuracy, macro-F1 (classes with support), weighted-F1 and per-class stats."""
|
||||
import numpy as np
|
||||
|
||||
matrix = confusion_matrix(y_true, y_pred, num_classes)
|
||||
total = int(matrix.sum())
|
||||
accuracy = float(np.trace(matrix) / total) if total else 0.0
|
||||
|
||||
per_class = []
|
||||
supports = []
|
||||
for c in range(num_classes):
|
||||
precision, recall, f1, support = _prf_from_confusion(matrix, c)
|
||||
per_class.append(
|
||||
{
|
||||
"class_id": c,
|
||||
"precision": precision,
|
||||
"recall": recall,
|
||||
"f1": f1,
|
||||
"support": support,
|
||||
}
|
||||
)
|
||||
supports.append(support)
|
||||
|
||||
supports = np.asarray(supports, dtype=np.float64)
|
||||
f1s = np.asarray([row["f1"] for row in per_class], dtype=np.float64)
|
||||
present = supports > 0
|
||||
macro_f1 = float(f1s[present].mean()) if present.any() else 0.0
|
||||
weighted_f1 = (
|
||||
float((f1s * supports).sum() / supports.sum()) if supports.sum() else 0.0
|
||||
)
|
||||
return {
|
||||
"accuracy": accuracy,
|
||||
"macro_f1": macro_f1,
|
||||
"weighted_f1": weighted_f1,
|
||||
"n_samples": total,
|
||||
"per_class": per_class,
|
||||
"confusion_matrix": matrix.astype(int).tolist(),
|
||||
}
|
||||
|
||||
|
||||
def rare_class_metrics(
|
||||
y_true,
|
||||
y_pred,
|
||||
num_classes: int,
|
||||
rare_class_ids: Sequence[int],
|
||||
) -> dict:
|
||||
"""Macro-F1 and recall restricted to *rare_class_ids*.
|
||||
|
||||
Classes with no evaluation samples are skipped. Returns the mean over the
|
||||
evaluated rare classes together with their support.
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
matrix = confusion_matrix(y_true, y_pred, num_classes)
|
||||
f1s, recalls, evaluated, supports = [], [], [], []
|
||||
for c in sorted(set(int(v) for v in rare_class_ids)):
|
||||
_, recall, f1, support = _prf_from_confusion(matrix, c)
|
||||
if support == 0:
|
||||
continue
|
||||
f1s.append(f1)
|
||||
recalls.append(recall)
|
||||
supports.append(support)
|
||||
evaluated.append(c)
|
||||
return {
|
||||
"rare_class_macro_f1": float(np.mean(f1s)) if f1s else 0.0,
|
||||
"rare_class_recall": float(np.mean(recalls)) if recalls else 0.0,
|
||||
"evaluated_rare_classes": evaluated,
|
||||
"rare_support": int(sum(supports)),
|
||||
}
|
||||
|
||||
|
||||
def unseen_class_report(y_true, y_pred, num_classes: int, unseen_class_ids: Sequence[int]) -> dict:
|
||||
"""Recall over samples whose true class was unseen during training."""
|
||||
import numpy as np
|
||||
|
||||
matrix = confusion_matrix(y_true, y_pred, num_classes)
|
||||
rows = []
|
||||
for c in sorted(set(int(v) for v in unseen_class_ids)):
|
||||
_, recall, f1, support = _prf_from_confusion(matrix, c)
|
||||
if support == 0:
|
||||
continue
|
||||
rows.append(
|
||||
{"class_id": c, "recall": recall, "f1": f1, "support": support}
|
||||
)
|
||||
return {"unseen_class_rows": rows}
|
||||
|
||||
|
||||
def mean_metric(per_client: Sequence[Mapping[str, float]], key: str) -> float:
|
||||
values = [float(m[key]) for m in per_client if key in m]
|
||||
return float(sum(values) / len(values)) if values else 0.0
|
||||
@@ -0,0 +1,15 @@
|
||||
"""Federated client updates, server aggregation and round orchestration."""
|
||||
|
||||
from .client import ClientUpdate, local_update
|
||||
from .round import TrainingConfig, run_federated
|
||||
from .server import aggregate_models, aggregate_prototypes, predict_global
|
||||
|
||||
__all__ = [
|
||||
"ClientUpdate",
|
||||
"local_update",
|
||||
"TrainingConfig",
|
||||
"run_federated",
|
||||
"aggregate_models",
|
||||
"aggregate_prototypes",
|
||||
"predict_global",
|
||||
]
|
||||
@@ -0,0 +1,130 @@
|
||||
"""Local client training (baseline PROTEAN objective, frozen).
|
||||
|
||||
Loss (paper Eq. 3 / thesis Eq. 5.6)::
|
||||
|
||||
CE(logits, y) + lambda * ||phi(x) - Cbar_y||^2
|
||||
+ (mu / 2) * ||omega_i - omega||^2
|
||||
|
||||
The prototype-alignment term is applied per sample for classes that currently
|
||||
have a global prototype (FedProto-style reading of Eq. 3); at round 0 no global
|
||||
prototypes exist and the term is skipped. Local prototypes are recomputed at the
|
||||
end of the local round as mean embeddings over the client's class samples
|
||||
(paper Eq. 1).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Mapping, Optional
|
||||
|
||||
from ..models.protean_cnn import get_parameters, set_parameters
|
||||
from ..models.prototypes import compute_local_prototypes
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClientUpdate:
|
||||
client_id: int
|
||||
parameters: list
|
||||
prototypes: dict = field(default_factory=dict)
|
||||
supports: dict = field(default_factory=dict)
|
||||
n_samples: int = 0
|
||||
loss: float = 0.0
|
||||
|
||||
|
||||
def local_update(
|
||||
client_id: int,
|
||||
model,
|
||||
global_parameters,
|
||||
global_prototypes: Optional[Mapping[int, object]],
|
||||
features,
|
||||
labels,
|
||||
indices,
|
||||
*,
|
||||
num_classes: int,
|
||||
local_epochs: int,
|
||||
batch_size: int,
|
||||
learning_rate: float,
|
||||
mu: float,
|
||||
lam: float,
|
||||
seed: int,
|
||||
) -> ClientUpdate:
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
if len(indices) == 0:
|
||||
return ClientUpdate(
|
||||
client_id=client_id,
|
||||
parameters=[
|
||||
torch.as_tensor(p, dtype=p.dtype).clone() for p in global_parameters
|
||||
],
|
||||
prototypes={},
|
||||
supports={},
|
||||
n_samples=0,
|
||||
)
|
||||
|
||||
device = next(model.parameters()).device
|
||||
set_parameters(model, global_parameters)
|
||||
|
||||
x = torch.as_tensor(features[indices], dtype=torch.float32, device=device)
|
||||
y = torch.as_tensor(np.asarray(labels)[indices], 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)
|
||||
n = x.shape[0]
|
||||
last_loss = 0.0
|
||||
|
||||
model.train()
|
||||
for _ in range(local_epochs):
|
||||
order = torch.randperm(n, generator=generator)
|
||||
for start in range(0, n, batch_size):
|
||||
batch_idx = order[start : start + batch_size]
|
||||
xb = x[batch_idx]
|
||||
yb = y[batch_idx]
|
||||
optimizer.zero_grad()
|
||||
log_prob = model(xb)
|
||||
loss = torch.nn.functional.nll_loss(log_prob, yb)
|
||||
|
||||
if proto_tensors:
|
||||
embeddings = 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 = embeddings[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()
|
||||
last_loss = float(loss.detach().cpu().item())
|
||||
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
embeddings = model.embedding(x).cpu().numpy()
|
||||
prototypes, supports = compute_local_prototypes(
|
||||
embeddings, np.asarray(labels)[indices], num_classes
|
||||
)
|
||||
return ClientUpdate(
|
||||
client_id=client_id,
|
||||
parameters=get_parameters(model),
|
||||
prototypes=prototypes,
|
||||
supports=supports,
|
||||
n_samples=int(n),
|
||||
loss=last_loss,
|
||||
)
|
||||
@@ -0,0 +1,158 @@
|
||||
"""Federated round orchestration for the baseline PROTEAN pipeline.
|
||||
|
||||
Per round: broadcast the global model and global prototypes, run local updates,
|
||||
average model parameters (1/M), aggregate class prototypes over ``S_j``, then
|
||||
evaluate the updated global prototypes on the validation split. The test split is
|
||||
never touched here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Mapping, Optional, Sequence
|
||||
|
||||
from ..evaluation.metrics import classification_metrics, rare_class_metrics
|
||||
from ..models.protean_cnn import (
|
||||
ModelConfig,
|
||||
build_model,
|
||||
extract_embeddings,
|
||||
get_parameters,
|
||||
set_parameters,
|
||||
)
|
||||
from ..models.prototypes import nearest_prototype_predict
|
||||
from ..utils.seeding import derive_seed, seed_everything
|
||||
from . import client as client_mod
|
||||
from . import server as server_mod
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainingConfig:
|
||||
rounds: int = 10
|
||||
local_epochs: int = 3
|
||||
batch_size: int = 256
|
||||
learning_rate: float = 1e-3
|
||||
mu: float = 0.1
|
||||
lam: float = 1.0
|
||||
num_clients: int = 10
|
||||
embedding_dim: int = 128
|
||||
dense_units: int = 128
|
||||
snapshot_round: Optional[int] = 0
|
||||
device: str = "cpu"
|
||||
|
||||
|
||||
def _evaluate(model, features, labels, global_prototypes, num_classes, rare_class_ids=None):
|
||||
if not global_prototypes:
|
||||
return {"accuracy": 0.0, "macro_f1": 0.0, "n_samples": 0}
|
||||
embeddings = extract_embeddings(model, features)
|
||||
predictions = nearest_prototype_predict(embeddings, global_prototypes)
|
||||
metrics = classification_metrics(labels, predictions, num_classes)
|
||||
summary = {
|
||||
"accuracy": metrics["accuracy"],
|
||||
"macro_f1": metrics["macro_f1"],
|
||||
"n_samples": metrics["n_samples"],
|
||||
}
|
||||
if rare_class_ids:
|
||||
rare = rare_class_metrics(labels, predictions, num_classes, rare_class_ids)
|
||||
summary["rare_class_macro_f1"] = rare["rare_class_macro_f1"]
|
||||
summary["rare_class_recall"] = rare["rare_class_recall"]
|
||||
return summary
|
||||
|
||||
|
||||
def run_federated(
|
||||
*,
|
||||
train_features,
|
||||
train_labels,
|
||||
assignments: Sequence,
|
||||
val_features,
|
||||
val_labels,
|
||||
num_classes: int,
|
||||
input_dim: int,
|
||||
config: TrainingConfig,
|
||||
init_seed_parts: Sequence[object],
|
||||
rare_class_ids: Optional[Sequence[int]] = None,
|
||||
) -> dict:
|
||||
"""Run the baseline federated loop and return fit artifacts and history."""
|
||||
import numpy as np
|
||||
|
||||
seed_everything(*init_seed_parts, "model_init")
|
||||
model_cfg = ModelConfig(
|
||||
input_dim=int(input_dim),
|
||||
num_classes=int(num_classes),
|
||||
embedding_dim=config.embedding_dim,
|
||||
dense_units=config.dense_units,
|
||||
)
|
||||
model = build_model(model_cfg)
|
||||
model.to(config.device)
|
||||
|
||||
client_models = []
|
||||
for _ in range(config.num_clients):
|
||||
cm = build_model(model_cfg)
|
||||
cm.to(config.device)
|
||||
client_models.append(cm)
|
||||
|
||||
global_parameters = get_parameters(model)
|
||||
global_prototypes: dict[int, object] = {}
|
||||
history: list[dict] = []
|
||||
snapshot: dict | None = None
|
||||
|
||||
for round_id in range(config.rounds):
|
||||
updates = []
|
||||
for client_id in range(config.num_clients):
|
||||
update = client_mod.local_update(
|
||||
client_id,
|
||||
client_models[client_id],
|
||||
global_parameters,
|
||||
global_prototypes,
|
||||
train_features,
|
||||
train_labels,
|
||||
assignments[client_id],
|
||||
num_classes=num_classes,
|
||||
local_epochs=config.local_epochs,
|
||||
batch_size=config.batch_size,
|
||||
learning_rate=config.learning_rate,
|
||||
mu=config.mu,
|
||||
lam=config.lam,
|
||||
seed=derive_seed(*init_seed_parts, "client_shuffle", round_id, client_id),
|
||||
)
|
||||
updates.append(update)
|
||||
|
||||
global_parameters = server_mod.aggregate_models([u.parameters for u in updates])
|
||||
set_parameters(model, global_parameters)
|
||||
global_prototypes = server_mod.aggregate_prototypes(
|
||||
[u.prototypes for u in updates],
|
||||
[u.supports for u in updates],
|
||||
num_classes,
|
||||
)
|
||||
|
||||
if config.snapshot_round is not None and round_id == config.snapshot_round:
|
||||
snapshot = {
|
||||
"round": round_id,
|
||||
"client_prototypes": [
|
||||
{int(c): np.asarray(v).tolist() for c, v in u.prototypes.items()}
|
||||
for u in updates
|
||||
],
|
||||
"client_supports": [
|
||||
{int(c): int(n) for c, n in u.supports.items()} for u in updates
|
||||
],
|
||||
}
|
||||
|
||||
row = {"round": round_id, "n_global_prototypes": len(global_prototypes)}
|
||||
row.update(
|
||||
_evaluate(
|
||||
model,
|
||||
val_features,
|
||||
val_labels,
|
||||
global_prototypes,
|
||||
num_classes,
|
||||
rare_class_ids,
|
||||
)
|
||||
)
|
||||
history.append(row)
|
||||
|
||||
return {
|
||||
"model": model,
|
||||
"parameters": global_parameters,
|
||||
"global_prototypes": global_prototypes,
|
||||
"history": history,
|
||||
"snapshot": snapshot,
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
"""Server-side aggregation and global inference (baseline PROTEAN).
|
||||
|
||||
Model parameters are averaged with equal ``1/M`` weights (paper Eq. 2/5.5).
|
||||
Class prototypes are aggregated by the active mode over the contributor set
|
||||
``S_j``; missing classes are never fabricated. Inference assigns the class of the
|
||||
nearest global prototype by squared L2 distance (thesis Eqs. 5.11-5.13).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Mapping, Sequence
|
||||
|
||||
from ..aggregation import AggregationMode
|
||||
from ..models.protean_cnn import average_parameters
|
||||
from ..models.prototypes import collect_global_prototypes
|
||||
|
||||
|
||||
def aggregate_models(client_parameter_lists: Sequence[list]) -> list:
|
||||
"""Equal-weight (1/M) federated averaging of client model parameters."""
|
||||
return average_parameters(list(client_parameter_lists))
|
||||
|
||||
|
||||
def aggregate_prototypes(
|
||||
client_prototypes: Sequence[Mapping[int, object]],
|
||||
client_supports: Sequence[Mapping[int, int]],
|
||||
num_classes: int,
|
||||
mode: AggregationMode | str = AggregationMode.ORIGINAL_PROTEAN,
|
||||
):
|
||||
"""Aggregate client prototypes into ``{class_id: Cbar_j}`` (baseline)."""
|
||||
return collect_global_prototypes(
|
||||
client_prototypes, client_supports, num_classes, mode=mode
|
||||
)
|
||||
|
||||
|
||||
def predict_global(model, features, global_prototypes):
|
||||
"""Predict labels via nearest global prototype (squared L2)."""
|
||||
from ..models.protean_cnn import extract_embeddings
|
||||
from ..models.prototypes import nearest_prototype_predict
|
||||
|
||||
if not global_prototypes:
|
||||
raise ValueError(
|
||||
"no global prototypes available: cannot perform prototype inference"
|
||||
)
|
||||
embeddings = extract_embeddings(model, features)
|
||||
return nearest_prototype_predict(embeddings, global_prototypes)
|
||||
@@ -0,0 +1,27 @@
|
||||
"""Baseline PROTEAN model and prototype utilities."""
|
||||
|
||||
from .protean_cnn import (
|
||||
ModelConfig,
|
||||
average_parameters,
|
||||
build_model,
|
||||
extract_embeddings,
|
||||
get_parameters,
|
||||
set_parameters,
|
||||
)
|
||||
from .prototypes import (
|
||||
collect_global_prototypes,
|
||||
compute_local_prototypes,
|
||||
nearest_prototype_predict,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"ModelConfig",
|
||||
"build_model",
|
||||
"get_parameters",
|
||||
"set_parameters",
|
||||
"average_parameters",
|
||||
"extract_embeddings",
|
||||
"compute_local_prototypes",
|
||||
"collect_global_prototypes",
|
||||
"nearest_prototype_predict",
|
||||
]
|
||||
@@ -0,0 +1,116 @@
|
||||
"""The PROTEAN detection CNN (baseline architecture, frozen).
|
||||
|
||||
Two Conv1D blocks (64 then 128 filters, ReLU, max-pool, dropout 0.2/0.5), a
|
||||
flatten, a dense embedding layer (dimension 128) and a log-softmax
|
||||
classification head. Tabular flow features are treated as a single-channel 1D
|
||||
input, as the baseline config describes. The architecture is not part of this
|
||||
study's contribution and must not be altered.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelConfig:
|
||||
input_dim: int
|
||||
num_classes: int
|
||||
embedding_dim: int = 128
|
||||
conv1_filters: int = 64
|
||||
conv2_filters: int = 128
|
||||
kernel_size: int = 3
|
||||
dropout1: float = 0.2
|
||||
dropout2: float = 0.5
|
||||
dense_units: int = 128
|
||||
|
||||
|
||||
def build_model(config: ModelConfig):
|
||||
"""Instantiate the CNN from a :class:`ModelConfig`."""
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class ProteanCNN(nn.Module):
|
||||
def __init__(self, cfg: ModelConfig) -> None:
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
self.conv1 = nn.Conv1d(1, cfg.conv1_filters, kernel_size=cfg.kernel_size, padding=cfg.kernel_size // 2)
|
||||
self.conv2 = nn.Conv1d(cfg.conv1_filters, cfg.conv2_filters, kernel_size=cfg.kernel_size, padding=cfg.kernel_size // 2)
|
||||
self.pool = nn.MaxPool1d(2)
|
||||
self.drop1 = nn.Dropout(cfg.dropout1)
|
||||
self.drop2 = nn.Dropout(cfg.dropout2)
|
||||
self.flatten = nn.Flatten()
|
||||
with torch.no_grad():
|
||||
dummy = torch.zeros(1, 1, cfg.input_dim)
|
||||
out = self._conv_blocks(dummy)
|
||||
flat_dim = int(out.shape[1])
|
||||
self.fc = nn.Linear(flat_dim, cfg.dense_units)
|
||||
self.head = nn.Linear(cfg.dense_units, cfg.num_classes)
|
||||
|
||||
def _conv_blocks(self, x):
|
||||
import torch.nn.functional as F
|
||||
|
||||
h = F.relu(self.conv1(x))
|
||||
h = self.pool(h)
|
||||
h = self.drop1(h)
|
||||
h = F.relu(self.conv2(h))
|
||||
h = self.pool(h)
|
||||
h = self.drop2(h)
|
||||
return self.flatten(h)
|
||||
|
||||
def embedding(self, features):
|
||||
"""Return phi(x), the dense embedding used for prototypes."""
|
||||
x = features.unsqueeze(1)
|
||||
h = self._conv_blocks(x)
|
||||
return self.fc(h)
|
||||
|
||||
def forward(self, features):
|
||||
import torch.nn.functional as F
|
||||
|
||||
return F.log_softmax(self.head(self.embedding(features)), dim=-1)
|
||||
|
||||
return ProteanCNN(config)
|
||||
|
||||
|
||||
def get_parameters(model) -> list:
|
||||
"""Return a deep copy of the model parameters as detached tensors."""
|
||||
return [p.detach().clone() for p in model.parameters()]
|
||||
|
||||
|
||||
def set_parameters(model, parameters) -> None:
|
||||
"""Load a parameter list (same shapes/order as ``model.parameters()``)."""
|
||||
import torch
|
||||
|
||||
with torch.no_grad():
|
||||
for param, new in zip(model.parameters(), parameters):
|
||||
param.copy_(torch.as_tensor(new, dtype=param.dtype, device=param.device))
|
||||
|
||||
|
||||
def average_parameters(param_lists: list[list]) -> list:
|
||||
"""Equal-weight (1/M) average of client parameter lists (paper Eq. 2/5.5)."""
|
||||
import torch
|
||||
|
||||
if not param_lists:
|
||||
raise ValueError("param_lists must not be empty")
|
||||
count = len(param_lists)
|
||||
averaged = []
|
||||
for tensors in zip(*param_lists):
|
||||
stacked = torch.stack([t.detach().float() for t in tensors], dim=0)
|
||||
averaged.append(stacked.mean(dim=0))
|
||||
return averaged
|
||||
|
||||
|
||||
def extract_embeddings(model, features, batch_size: int = 1024):
|
||||
"""Return ``phi(x)`` as a NumPy array, in eval mode, without gradients."""
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
model.eval()
|
||||
outputs: list[np.ndarray] = []
|
||||
with torch.no_grad():
|
||||
for start in range(0, len(features), batch_size):
|
||||
batch = torch.as_tensor(features[start : start + batch_size], dtype=torch.float32)
|
||||
outputs.append(model.embedding(batch).cpu().numpy())
|
||||
if not outputs:
|
||||
return np.empty((0, model.cfg.dense_units), dtype=np.float32)
|
||||
return np.concatenate(outputs, axis=0).astype(np.float32)
|
||||
@@ -0,0 +1,78 @@
|
||||
"""Local prototype construction and global prototype collection.
|
||||
|
||||
Baseline semantics (paper Eq. 1 / thesis Eq. 5.4): the local prototype ``C_ij``
|
||||
is the mean embedding over the client's class-``j`` training samples. Missing
|
||||
classes (``N_ij = 0``) produce no prototype; missing classes are never
|
||||
fabricated. Global prototypes are built by the server from the contributor set
|
||||
``S_j`` via the aggregation module.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Mapping, Sequence
|
||||
|
||||
from ..aggregation import AggregationMode, aggregate_class_prototypes
|
||||
|
||||
|
||||
def compute_local_prototypes(embeddings, labels, num_classes: int):
|
||||
"""Return ``(prototypes, supports)`` for the classes present on a client.
|
||||
|
||||
``prototypes[c]`` is an ``np.ndarray`` mean embedding; ``supports[c]`` is the
|
||||
sample count. Absent classes are omitted (no zero vectors are emitted).
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
embeddings = np.asarray(embeddings)
|
||||
labels = np.asarray(labels)
|
||||
prototypes: dict[int, "np.ndarray"] = {}
|
||||
supports: dict[int, int] = {}
|
||||
for class_id in range(num_classes):
|
||||
mask = labels == class_id
|
||||
n = int(mask.sum())
|
||||
if n == 0:
|
||||
continue
|
||||
prototypes[class_id] = embeddings[mask].mean(axis=0).astype(np.float32)
|
||||
supports[class_id] = n
|
||||
return prototypes, supports
|
||||
|
||||
|
||||
def collect_global_prototypes(
|
||||
client_prototypes: Sequence[Mapping[int, object]],
|
||||
client_supports: Sequence[Mapping[int, int]],
|
||||
num_classes: int,
|
||||
mode: AggregationMode | str = AggregationMode.ORIGINAL_PROTEAN,
|
||||
):
|
||||
"""Server-side aggregation of per-client prototypes into global prototypes.
|
||||
|
||||
Clients that did not observe a class simply do not appear in
|
||||
``client_prototypes`` for that class, so ``S_j`` is built implicitly and
|
||||
missing classes are never inserted as zero vectors.
|
||||
"""
|
||||
contributions: dict[int, list] = {c: [] for c in range(num_classes)}
|
||||
for prototypes, supports in zip(client_prototypes, client_supports):
|
||||
for class_id, proto in prototypes.items():
|
||||
contributions[int(class_id)].append((proto, supports.get(int(class_id), 0)))
|
||||
return aggregate_class_prototypes(contributions, mode)
|
||||
|
||||
|
||||
def nearest_prototype_predict(embeddings, global_prototypes: Mapping[int, object]):
|
||||
"""Predict by minimum squared L2 distance to global prototypes (Eq. 5.11-5.13).
|
||||
|
||||
``global_prototypes`` maps class id -> prototype. Samples whose embedding has
|
||||
no prototype to compare against (empty prototype set) raise ``ValueError``.
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
if not global_prototypes:
|
||||
raise ValueError("no global prototypes available for inference")
|
||||
class_ids = sorted(global_prototypes)
|
||||
matrix = np.stack([np.asarray(global_prototypes[c], dtype=np.float64) for c in class_ids])
|
||||
embeddings = np.asarray(embeddings, dtype=np.float64)
|
||||
# Squared L2 via the expanded norm identity.
|
||||
dists = (
|
||||
(embeddings ** 2).sum(axis=1, keepdims=True)
|
||||
- 2.0 * embeddings @ matrix.T
|
||||
+ (matrix ** 2).sum(axis=1)[None, :]
|
||||
)
|
||||
chosen = dists.argmin(axis=1)
|
||||
return np.asarray([class_ids[i] for i in chosen], dtype=np.int64)
|
||||
@@ -0,0 +1,39 @@
|
||||
"""Integrity helpers: file and content hashing for reproducible manifests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
|
||||
def sha256_bytes(data: bytes) -> str:
|
||||
return hashlib.sha256(data).hexdigest()
|
||||
|
||||
|
||||
def sha256_file(path: str | Path, chunk_size: int = 1 << 20) -> str:
|
||||
"""Stream-hash a file without loading it fully into memory."""
|
||||
digest = hashlib.sha256()
|
||||
with Path(path).open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(chunk_size), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def hash_json(obj: Any) -> str:
|
||||
"""Stable SHA-256 of a JSON-serialisable object (sorted keys)."""
|
||||
payload = json.dumps(obj, sort_keys=True, separators=(",", ":"), default=str)
|
||||
return sha256_bytes(payload.encode("utf-8"))
|
||||
|
||||
|
||||
def hash_index_arrays(arrays: Iterable[object]) -> str:
|
||||
"""Stable hash of a sequence of integer index arrays.
|
||||
|
||||
Arrays are normalised to sorted Python int lists so the hash does not depend
|
||||
on NumPy dtype or memory layout.
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
normalised = [sorted(int(v) for v in np.asarray(a).ravel()) for a in arrays]
|
||||
return hash_json(normalised)
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Deterministic seeding helpers.
|
||||
|
||||
All randomness in the project must be reproducible and isolated by purpose, so
|
||||
that two aggregation variants sharing ``(dataset, alpha, seed)`` also share the
|
||||
split, partition, model initialisation and batch order.
|
||||
|
||||
Use :func:`derive_seed` to turn semantic identifiers into a stable integer seed,
|
||||
never ``hash()`` (which is salted per process).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import random
|
||||
|
||||
|
||||
def derive_seed(*parts: object, mod: int = 2**31) -> int:
|
||||
"""Return a stable non-negative integer seed derived from *parts*.
|
||||
|
||||
The result depends only on the string form of *parts*, so it is stable
|
||||
across processes, machines and Python versions.
|
||||
"""
|
||||
payload = "|".join(str(p) for p in parts).encode("utf-8")
|
||||
digest = hashlib.sha256(payload).hexdigest()
|
||||
return int(digest[:16], 16) % mod
|
||||
|
||||
|
||||
def numpy_rng(*parts: object):
|
||||
"""Return a seeded :class:`numpy.random.Generator` for *parts*."""
|
||||
import numpy as np
|
||||
|
||||
return np.random.default_rng(derive_seed(*parts))
|
||||
|
||||
|
||||
def seed_everything(*parts: object) -> int:
|
||||
"""Seed Python, NumPy and PyTorch RNGs deterministically from *parts*.
|
||||
|
||||
Returns the derived integer seed for recording in run metadata.
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
seed = derive_seed(*parts)
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
try:
|
||||
import torch
|
||||
|
||||
torch.manual_seed(seed)
|
||||
except ImportError: # pragma: no cover - torch is a research dependency
|
||||
pass
|
||||
return seed
|
||||
Reference in new issue
Block a user