"""P076: detect broken cross-field relationships using hidden-field reconstruction."""
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 = "P076"
BUDGET = 500
GROUPS = ("amber", "birch", "cove", "dune")
SLOPES = np.array([0.7, -0.8, 1.2, -1.3])
OFFSETS = np.array([-0.5, 0.5, -0.3, 0.3])
TRAIN_ROWS = 4096
TEST_ROWS = 2048
def records(*, rows: int, seed: int, shuffled: bool = False, tail: bool = False) -> Iterator[dict]:
rng = np.random.default_rng(seed)
group = np.resize(np.arange(len(GROUPS)), rows)
rng.shuffle(group)
x = rng.uniform(-2, 2, rows)
if tail:
x = rng.choice([-1.0, 1.0], rows) * rng.uniform(1.8, 2.0, rows)
y = SLOPES[group] * x + OFFSETS[group] + rng.normal(0, 0.1, rows)
if shuffled:
y = y[rng.permutation(rows)]
for index in range(rows):
yield {"group": GROUPS[group[index]], "x": float(x[index]), "y": float(y[index])}
def build() -> rf.Model:
model = rf.Model.xs(
batch_size=128,
group=rf.Enum(values=GROUPS, p_unavailable=0.0),
x=rf.Number,
y=rf.Number(mask=True, objective="mse"),
)
model.optimizer = rf.adamw(learning_rate=0.002)
return model
def predict(model: rf.Model, rows: list[dict]) -> np.ndarray:
visible = [{"group": row["group"], "x": row["x"]} for row in rows]
result = model.predict(pa.Table.from_pylist(visible))["predictions"].to_pylist()
return np.asarray([row["y"]["content"] for row in result], dtype=float)
def auc(normal: np.ndarray, anomalous: np.ndarray) -> float:
"""Probability that an anomaly scores above a normal row, with half credit for ties."""
ordered = np.sort(normal)
left = np.searchsorted(ordered, anomalous, side="left")
right = np.searchsorted(ordered, anomalous, side="right")
return float(np.mean((left + right) / (2 * len(normal))))
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
stream = seed * 100
validation = list(records(rows=2048, seed=stream + 2))
test = list(records(rows=TEST_ROWS, seed=stream + 3))
anomalous = list(records(rows=TEST_ROWS, seed=stream + 3, shuffled=True))
tail = list(records(rows=TEST_ROWS, seed=stream + 4, tail=True))
actual = np.asarray([row["y"] for row in test])
broken = np.asarray([row["y"] for row in anomalous])
train_mean = float(np.mean([row["y"] for row in records(rows=TRAIN_ROWS, seed=stream + 1)]))
baseline_rmse = float(np.sqrt(np.mean((actual - train_mean) ** 2)))
marginal_auc = auc(np.abs(actual - train_mean), np.abs(broken - train_mean))
metrics = {"train_rows": TRAIN_ROWS, "test_rows_per_condition": TEST_ROWS, "marginal_auc": marginal_auc}
checks = {
"Anomalies preserve the exact target marginal": bool(np.array_equal(np.sort(actual), np.sort(broken))),
"Anomalies preserve visible contexts": all(
(a["group"], a["x"]) == (b["group"], b["x"]) for a, b in zip(test, anomalous, strict=True)
),
"Marginal deviation baseline has exactly chance AUROC": abs(marginal_auc - 0.5) < 1e-12,
}
for shuffled in (False, True):
lit.seed_everything(seed, workers=True)
model = build()
data = rf.SyntheticDataModule(
model=model,
train=partial(records, rows=TRAIN_ROWS, seed=stream + 1, shuffled=shuffled),
validate=partial(records, rows=2048, seed=stream + 2, shuffled=shuffled),
seed=seed,
)
trainer = lit.Trainer(
accelerator=accelerator,
devices=1,
max_epochs=-1,
max_steps=BUDGET if steps is None else min(steps, BUDGET),
deterministic=True,
logger=False,
enable_checkpointing=False,
enable_progress_bar=False,
enable_model_summary=False,
num_sanity_val_steps=0,
)
trainer.fit(model, datamodule=data)
model.eval()
valid_error = np.abs(predict(model, validation) - [row["y"] for row in validation])
threshold = float(np.quantile(valid_error, 0.95))
prediction = predict(model, test)
clean_score = np.abs(prediction - actual)
broken_score = np.abs(prediction - broken)
tail_score = np.abs(predict(model, tail) - [row["y"] for row in tail])
key = "shuffled_training" if shuffled else "informative"
result = {
"steps": trainer.global_step,
"validation_threshold": threshold,
"rmse": float(np.sqrt(np.mean(clean_score**2))),
"baseline_rmse": baseline_rmse,
"normalized_rmse": float(np.sqrt(np.mean(clean_score**2)) / baseline_rmse),
"anomaly_auc": auc(clean_score, broken_score),
"detection_recall": float(np.mean(broken_score > threshold)),
"normal_false_positive_rate": float(np.mean(clean_score > threshold)),
"tail_false_positive_rate": float(np.mean(tail_score > threshold)),
}
metrics[key] = result
checks[f"{key}: every measurement is finite"] = bool(np.isfinite(list(result.values())).all())
if shuffled:
checks["Shuffled training anomaly AUROC stays below 0.65"] = result["anomaly_auc"] <= 0.65
else:
checks["Informative reconstruction nRMSE below 0.20"] = result["normalized_rmse"] < 0.20
checks["Contextual anomaly AUROC reaches 0.90"] = result["anomaly_auc"] >= 0.90
checks["Anomaly recall reaches 0.75"] = result["detection_recall"] >= 0.75
checks["Normal false positive rate at most 0.10"] = result["normal_false_positive_rate"] <= 0.10
checks["Rare valid tail false positive rate at most 0.15"] = result["tail_false_positive_rate"] <= 0.15
return metrics, checks
if __name__ == "__main__":
report(PROOF_ID, run, seed=7601)