"""P078: require real partition recovery and held-out skill from Cluster."""
from collections.abc import Iterator
from functools import partial
import lightning.pytorch as lit
import numpy as np
import pyarrow as pa
from reporting import report
import relflow as rf
PROOF_ID = "P078"
BUDGET = 600
BATCH_SIZE = 128
TRAIN_ROWS = 4096
TEST_ROWS = 2048
ENTITY_COUNT = 64
REGIME_COUNT = 4
INTERCEPTS = np.asarray((-6.0, -2.0, 2.0, 6.0))
SLOPES = np.asarray((-0.6, 0.6, -0.6, 0.6))
def partition(seed: int) -> np.ndarray:
"""A balanced random mapping known only to the generator and evaluator."""
groups = np.arange(ENTITY_COUNT) % REGIME_COUNT
return np.random.default_rng(seed).permutation(groups)
def records(*, rows: int, seed: int, partition_seed: int, no_regime: bool = False) -> Iterator[dict]:
"""Preserve numeric records exactly when shuffling the stable identity link."""
rng = np.random.default_rng(seed)
entities = np.resize(np.arange(ENTITY_COUNT), rows)
rng.shuffle(entities)
groups = partition(partition_seed)[entities]
x = rng.uniform(-1.0, 1.0, rows)
y = INTERCEPTS[groups] + SLOPES[groups] * x + rng.normal(0.0, 0.1, rows)
if no_regime:
entities = entities[rng.permutation(rows)]
for entity, value, target in zip(entities, x, y, strict=True):
yield {"entity": f"entity-{entity:04d}", "x": float(value), "y": float(target)}
def adjusted_rand(actual: np.ndarray, predicted: np.ndarray) -> float:
"""Adjusted Rand agreement, invariant to the arbitrary cluster labels."""
_, left = np.unique(actual, return_inverse=True)
_, right = np.unique(predicted, return_inverse=True)
cells = np.zeros((left.max() + 1, right.max() + 1), dtype=np.int64)
np.add.at(cells, (left, right), 1)
joint = float(np.sum(cells * (cells - 1) / 2))
row = cells.sum(axis=1)
column = cells.sum(axis=0)
left_pairs = float(np.sum(row * (row - 1) / 2))
right_pairs = float(np.sum(column * (column - 1) / 2))
all_pairs = len(actual) * (len(actual) - 1) / 2
expected = left_pairs * right_pairs / all_pairs
maximum = (left_pairs + right_pairs) / 2
return 1.0 if maximum == expected else (joint - expected) / (maximum - expected)
def build() -> rf.Model:
model = rf.Model.xs(
batch_size=BATCH_SIZE,
entity=rf.Cluster(bounds=(2, 8), p_unavailable=0.0, mask=rf.Mask(rate=0.1, reconstruct=True)),
x=rf.Number,
y=rf.Number(mask=True, objective="mse"),
)
model.optimizer = rf.adamw(learning_rate=0.003)
return model
def predict(model: rf.Model, rows: list[dict]) -> np.ndarray:
visible = [{"entity": row["entity"], "x": row["x"]} for row in rows]
output = model.predict(pa.Table.from_pylist(visible))["predictions"].to_pylist()
return np.asarray([row["y"]["content"] for row in output])
def measure(model: rf.Model, rows: list[dict], training_mean: float, partition_seed: int) -> dict:
target = np.asarray([row["y"] for row in rows])
predicted = predict(model, rows)
baseline = float(np.sqrt(np.mean((target - training_mean) ** 2)))
rmse = float(np.sqrt(np.mean((target - predicted) ** 2)))
assignments = rf.Cluster.assignments(model, rf.Address("entity"))
groups = partition(partition_seed)
# Partial-budget smoke runs may not yet encounter every ID; absence remains
# an explicit failed coverage gate rather than silently dropping entities.
assigned = np.asarray(
[assignments.get(f"entity-{index:04d}", {}).get("cluster", -1) for index in range(ENTITY_COUNT)]
)
return {
"rmse": rmse,
"baseline_rmse": baseline,
"normalized_rmse": rmse / baseline,
"adjusted_rand": adjusted_rand(groups, assigned),
"assigned_entities": int((assigned >= 0).sum()),
"occupied_clusters": int(len(np.unique(assigned[assigned >= 0]))),
"assignments": assigned.tolist(),
"expected_partition": groups.tolist(),
"status": rf.Cluster.status(model, rf.Address("entity")),
}
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
metrics, checks = {}, {}
partition_seed = seed * 100
for arm, no_regime in (("stable", False), ("no_regime", True)):
lit.seed_everything(seed, workers=True)
model = build()
options = {"partition_seed": partition_seed, "no_regime": no_regime}
data = rf.SyntheticDataModule(
model=model,
train=partial(records, rows=TRAIN_ROWS, seed=seed * 100 + 1, **options),
validate=partial(records, rows=1024, seed=seed * 100 + 2, **options),
seed=seed,
)
trainer = lit.Trainer(
accelerator=accelerator,
devices=1,
max_steps=BUDGET if steps is None else min(steps, BUDGET),
max_epochs=-1,
logger=False,
enable_checkpointing=False,
enable_progress_bar=False,
enable_model_summary=False,
deterministic=True,
num_sanity_val_steps=0,
)
trainer.fit(model, datamodule=data)
model.eval()
state = rf.Cluster.assignments(model, rf.Address("entity"))
training_mean = float(np.mean([row["y"] for row in records(rows=TRAIN_ROWS, seed=seed * 100 + 1, **options)]))
test = list(records(rows=TEST_ROWS, seed=seed * 100 + 3, **options))
measured = measure(model, test, training_mean, partition_seed)
measured.update(steps=trainer.global_step, training_mean=training_mean, noise_rmse=0.1)
if arm == "stable":
shuffled = list(records(rows=TEST_ROWS, seed=seed * 100 + 3, partition_seed=partition_seed, no_regime=True))
measured["shuffled"] = measure(model, shuffled, training_mean, partition_seed)
metrics[arm] = measured
checks[f"{arm}: finite prediction and partition scores"] = bool(
np.isfinite([measured["rmse"], measured["normalized_rmse"], measured["adjusted_rand"]]).all()
)
checks[f"{arm}: all 64 known identities have assignments"] = measured["assigned_entities"] == ENTITY_COUNT
checks[f"{arm}: evaluation preserves assignment state"] = state == rf.Cluster.assignments(
model, rf.Address("entity")
)
stable, control = metrics["stable"], metrics["no_regime"]
checks.update(
{
"Stable regimes recover ARI at least 0.80": stable["adjusted_rand"] >= 0.80,
"Stable regimes predict with normalized RMSE at most 0.35": stable["normalized_rmse"] <= 0.35,
"No-regime control has ARI at most 0.15": control["adjusted_rand"] <= 0.15,
"No-regime control has normalized RMSE at least 0.85": control["normalized_rmse"] >= 0.85,
"Shuffling stable test IDs gives normalized RMSE at least 0.85": stable["shuffled"]["normalized_rmse"]
>= 0.85,
"Shuffling stable test IDs increases normalized RMSE by at least 0.40": (
stable["shuffled"]["normalized_rmse"] - stable["normalized_rmse"] >= 0.40
),
}
)
return metrics, checks
if __name__ == "__main__":
report(PROOF_ID, run, seed=7801)
