"""P075: test fresh-key retrieval with distractors and an intermediate relation."""
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 = "P075"
BUDGET = 600
BATCH_SIZE = 64
TRAIN_ROWS = 4096
TEST_ROWS = 2048
CAPACITIES = (2, 4, 8, 16)
def records(
*, rows: int, seed: int, namespace: str, hops: int, capacity: int | None = None, broken: bool = False
) -> Iterator[dict]:
"""Fresh unique keys with a matched key-rotation control and fixed targets."""
rng = np.random.default_rng(seed)
choices = CAPACITIES if hops == 1 else CAPACITIES[:2]
sizes = np.resize(choices, rows) if capacity is None else np.full(rows, capacity)
rng.shuffle(sizes)
for row, size in enumerate(sizes):
keys = [f"{namespace}-{row:06d}-s{index:02d}" for index in range(size)]
values = rng.uniform(-1.0, 1.0, size)
selected = rng.choice(size, size=2, replace=False)
source = [{"entity_id": keys[index], "value": float(values[index])} for index in rng.permutation(size)]
if hops == 1:
target = [
{"entity_id": keys[(index + int(broken)) % size], "value": float(values[index])} for index in selected
]
yield {"source": source, "target": target}
else:
query_keys = [f"{namespace}-{row:06d}-q{index:02d}" for index in range(size)]
links = [
{"entity_id": query_keys[index], "source_id": keys[(index + int(broken)) % size]}
for index in rng.permutation(size)
]
target = [{"entity_id": query_keys[index], "value": float(values[index])} for index in selected]
yield {"source": source, "links": links, "target": target}
def permute(rows: list[dict], seed: int) -> list[dict]:
"""Reorder complete records independently, preserving each visible relation."""
rng = np.random.default_rng(seed)
return [
{name: [items[index] for index in rng.permutation(len(items))] for name, items in row.items()} for row in rows
]
def build(hops: int) -> rf.Model:
fields = {
"source": rf.Branch(
length=16,
reduction=None,
n_layers=2,
entity_id=rf.Hash(n_hashes=4, n_bands=8),
value=rf.Number,
),
"target": rf.Branch(
length=2,
reduction=None,
n_layers=2,
entity_id=rf.Hash(n_hashes=4, n_bands=8),
value=rf.Number(mask=True, objective="mse"),
),
}
if hops == 2:
fields["links"] = rf.Branch(
length=16,
reduction=None,
n_layers=2,
entity_id=rf.Hash(n_hashes=4, n_bands=8),
source_id=rf.Hash(n_hashes=4, n_bands=8),
)
model = rf.Model.xs(n_layers=2, n_heads=4, reduction=None, batch_size=BATCH_SIZE, fields=fields)
model.optimizer = rf.adamw(learning_rate=0.003)
return model
def predict(model: rf.Model, rows: list[dict]) -> np.ndarray:
visible = [{**row, "target": [{"entity_id": item["entity_id"]} for item in row["target"]]} for row in rows]
output = model.predict(pa.Table.from_pylist(visible))["predictions"].to_pylist()
return np.asarray([[item["value"]["content"] for item in row["target"]] for row in output])
def score(predicted: np.ndarray, rows: list[dict], training_mean: float) -> dict:
actual = np.asarray([[item["value"] for item in row["target"]] for row in rows])
source_mean = np.asarray([np.mean([item["value"] for item in row["source"]]) for row in rows])[:, None]
baseline = float(np.sqrt(np.mean((actual - training_mean) ** 2)))
rmse = float(np.sqrt(np.mean((predicted - actual) ** 2)))
return {
"rmse": rmse,
"baseline_rmse": baseline,
"source_mean_rmse": float(np.sqrt(np.mean((actual - source_mean) ** 2))),
"normalized_rmse": rmse / baseline,
}
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
metrics, checks = {}, {}
for hops in (1, 2):
lit.seed_everything(seed, workers=True)
model = build(hops)
data = rf.SyntheticDataModule(
model=model,
train=partial(records, rows=TRAIN_ROWS, seed=seed * 100 + 1, namespace="train", hops=hops),
validate=partial(records, rows=512, seed=seed * 100 + 2, namespace="validate", hops=hops),
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()
train = records(rows=TRAIN_ROWS, seed=seed * 100 + 1, namespace="train", hops=hops)
training_mean = float(np.mean([item["value"] for row in train for item in row["target"]]))
metrics[f"hops_{hops}"] = {"steps": trainer.global_step, "training_mean": training_mean, "ladder": {}}
for capacity in CAPACITIES if hops == 1 else CAPACITIES[:2]:
options = {
"rows": TEST_ROWS,
"seed": seed * 100 + 10 + capacity,
"namespace": f"test-{hops}-{capacity}",
"hops": hops,
"capacity": capacity,
}
test = list(records(**options))
broken = list(records(**options, broken=True))
permuted = permute(test, seed * 100 + 40 + capacity)
intact = score(predict(model, test), test, training_mean)
corrupted = score(predict(model, broken), broken, training_mean)
reordered = score(predict(model, permuted), permuted, training_mean)
gap = corrupted["normalized_rmse"] - intact["normalized_rmse"]
delta = abs(reordered["normalized_rmse"] - intact["normalized_rmse"])
metrics[f"hops_{hops}"]["ladder"][capacity] = {
"intact": intact,
"broken": corrupted,
"permuted": reordered,
"corruption_gap": gap,
"permutation_nrmse_delta": delta,
}
label = f"{hops} hop, {capacity} sources"
checks.update(
{
f"{label}: finite source-unit errors": bool(
np.isfinite([*intact.values(), *corrupted.values(), *reordered.values()]).all()
),
f"{label}: intact nRMSE at most 0.50": intact["normalized_rmse"] <= 0.50,
f"{label}: beats source-mean RMSE by 20%": intact["rmse"] <= 0.8 * intact["source_mean_rmse"],
f"{label}: broken-key nRMSE at least 0.80": corrupted["normalized_rmse"] >= 0.80,
f"{label}: corruption increases nRMSE by at least 0.30": gap >= 0.30,
f"{label}: reordered nRMSE at most 0.50": reordered["normalized_rmse"] <= 0.50,
f"{label}: permutation nRMSE delta at most 0.10": delta <= 0.10,
}
)
return metrics, checks
if __name__ == "__main__":
report(PROOF_ID, run, seed=7501)
