"""P077: compare numerical representations and the quantities their objectives estimate."""
from collections.abc import Iterator
from functools import partial
import lightning.pytorch as lit
import numpy as np
import pyarrow as pa
import torch
from reporting import report
import relflow as rf
PROOF_ID = "P077"
BUDGET = 400
TRAIN_ROWS = 8192
TEST_ROWS = 4096
GROUPS = ("amber", "birch", "cove")
LOCATIONS = np.array([-0.5, 0.0, 0.5])
SIGMA = 0.8
def records(*, rows: int, seed: int, contaminated: bool = False) -> Iterator[dict]:
rng = np.random.default_rng(seed)
u = rng.uniform(0, 1, rows)
x = np.expm1(4 * u)
amount = 0.2 * x + rng.normal(0, 0.1, rows)
visible = x.copy()
if contaminated:
visible[rng.random(rows) < 0.02] *= 100
for index in range(rows):
yield {
"x": float(visible[index]),
"rank": "upper" if u[index] > 0.5 else "lower",
"amount": float(amount[index]),
}
def build(representation: str, task: str, objective: str = "mse") -> rf.Model:
scalar = rf.Quantile if representation == "quantile" else rf.Number
if task == "noisy":
model = rf.Model.xs(
batch_size=128,
group=rf.Enum(values=GROUPS, p_unavailable=0.0),
target=scalar(objective=objective, mask=True),
)
else:
target = (
rf.Enum(values=("lower", "upper"), p_unavailable=0.0, mask=True)
if task == "rank"
else rf.Number(objective="mse", mask=True)
)
model = rf.Model.xs(batch_size=128, x=scalar, fields={task: target})
model.optimizer = rf.adamw(learning_rate=0.002)
return model
def fit(model: rf.Model, train, validate, seed: int, steps: int | None, accelerator: str) -> int:
data = rf.SyntheticDataModule(model=model, train=train, validate=validate, 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()
return trainer.global_step
def predict(model: rf.Model, rows: list[dict], task: str) -> np.ndarray:
visible = [{"group": row["group"]} if task == "noisy" else {"x": row["x"]} for row in rows]
predictions = model.predict(pa.Table.from_pylist(visible))["predictions"].to_pylist()
name = "target" if task == "noisy" else task
content = [row[name]["content"] for row in predictions]
if task == "rank":
return np.asarray([row["value"] == "upper" for row in content])
return np.asarray(content, dtype=float)
def noisy_records(*, rows: int, seed: int) -> Iterator[dict]:
rng = np.random.default_rng(seed)
group = np.resize(np.arange(len(GROUPS)), rows)
rng.shuffle(group)
target = np.exp(LOCATIONS[group] + SIGMA * rng.standard_normal(rows))
for index in range(rows):
yield {"group": GROUPS[group[index]], "target": float(target[index])}
def cdf(values: np.ndarray) -> np.ndarray:
"""Analytical CDF of the equally weighted three-lognormal population mixture."""
scaled = (np.log(np.maximum(values, 1e-300))[..., None] - LOCATIONS) / SIGMA
return torch.special.ndtr(torch.from_numpy(scaled)).mean(dim=-1).numpy()
def inverse(percentiles: np.ndarray) -> np.ndarray:
lower = np.zeros_like(percentiles)
upper = np.full_like(percentiles, np.exp(LOCATIONS.max() + 10 * SIGMA))
for _ in range(64):
midpoint = (lower + upper) / 2
below = cdf(midpoint) < percentiles
lower = np.where(below, midpoint, lower)
upper = np.where(below, upper, midpoint)
return (lower + upper) / 2
def optima() -> dict[str, np.ndarray]:
differences = (LOCATIONS[:, None] - LOCATIONS[None, :]) / (np.sqrt(2) * SIGMA)
mean_percentiles = torch.special.ndtr(torch.from_numpy(differences)).mean(dim=1).numpy()
return {
"number_mse": np.exp(LOCATIONS + SIGMA**2 / 2),
"number_mae": np.exp(LOCATIONS),
"quantile_mse": inverse(mean_percentiles),
"quantile_mae": np.exp(LOCATIONS),
}
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
stream = seed * 100
test = list(records(rows=TEST_ROWS, seed=stream + 3))
tail = np.asarray([row["x"] > np.expm1(4 * 0.95) for row in test])
permutation = np.random.default_rng(stream + 4).permutation(TEST_ROWS)
shuffled = [{**row, "x": test[index]["x"]} for row, index in zip(test, permutation, strict=True)]
training = list(records(rows=TRAIN_ROWS, seed=stream + 1))
mean = float(np.mean([row["amount"] for row in training]))
target = np.asarray([row["amount"] for row in test])
baseline = float(np.sqrt(np.mean((target - mean) ** 2)))
labels = np.asarray([row["rank"] == "upper" for row in test])
metrics = {
"train_rows": TRAIN_ROWS,
"test_rows": TEST_ROWS,
"tail_rows": int(tail.sum()),
"central_rows": int((~tail).sum()),
"input_tasks": {},
"noisy_targets": {},
}
checks = {}
for task in ("rank", "amount"):
for representation in ("number", "quantile"):
for contaminated in (False, True):
lit.seed_everything(seed, workers=True)
model = build(representation, task)
updates = fit(
model,
partial(records, rows=TRAIN_ROWS, seed=stream + 1, contaminated=contaminated),
partial(records, rows=1024, seed=stream + 2),
seed,
steps,
accelerator,
)
output = predict(model, test, task)
control = predict(model, shuffled, task)
key = f"{task}_{representation}_{'contaminated' if contaminated else 'clean'}"
if task == "rank":
result = {
"steps": updates,
"accuracy": float(np.mean(output == labels)),
"central_accuracy": float(np.mean(output[~tail] == labels[~tail])),
"tail_accuracy": float(np.mean(output[tail] == labels[tail])),
"shuffled_accuracy": float(np.mean(control == labels)),
"majority_accuracy": float(max(labels.mean(), 1 - labels.mean())),
}
if not contaminated:
checks[f"{key}: accuracy reaches 0.90"] = result["accuracy"] >= 0.90
checks[f"{key}: shuffled accuracy at most 0.60"] = result["shuffled_accuracy"] <= 0.60
else:
result = {
"steps": updates,
"rmse": float(np.sqrt(np.mean((output - target) ** 2))),
"baseline_rmse": baseline,
"normalized_rmse": float(np.sqrt(np.mean((output - target) ** 2)) / baseline),
"tail_rmse": float(np.sqrt(np.mean((output[tail] - target[tail]) ** 2))),
"central_rmse": float(np.sqrt(np.mean((output[~tail] - target[~tail]) ** 2))),
"shuffled_normalized_rmse": float(np.sqrt(np.mean((control - target) ** 2)) / baseline),
}
if not contaminated:
checks[f"{key}: normalized RMSE at most 0.35"] = result["normalized_rmse"] <= 0.35
checks[f"{key}: shuffled normalized RMSE at least 0.90"] = (
result["shuffled_normalized_rmse"] >= 0.90
)
metrics["input_tasks"][key] = result
checks[f"{key}: finite measurements"] = bool(np.isfinite(list(result.values())).all())
noisy = list(noisy_records(rows=TEST_ROWS, seed=stream + 13))
groups = np.asarray([GROUPS.index(row["group"]) for row in noisy])
observed = np.asarray([row["target"] for row in noisy])
targets = optima()
metrics["population_optima"] = {key: values.tolist() for key, values in targets.items()}
for representation in ("number", "quantile"):
for objective in ("mse", "mae"):
lit.seed_everything(seed, workers=True)
model = build(representation, "noisy", objective)
updates = fit(
model,
partial(noisy_records, rows=TRAIN_ROWS, seed=stream + 11),
partial(noisy_records, rows=1024, seed=stream + 12),
seed,
steps,
accelerator,
)
output = predict(model, noisy, "noisy")
key = f"{representation}_{objective}"
optimum = targets[key][groups]
group_predictions = np.asarray([output[groups == group].mean() for group in range(len(GROUPS))])
relative = float(np.mean(np.abs(group_predictions - targets[key]) / targets[key]))
rank_error = cdf(output) - cdf(observed)
result = {
"steps": updates,
"group_predictions": group_predictions.tolist(),
"mean_relative_optimum_error": relative,
"source_rmse": float(np.sqrt(np.mean((output - observed) ** 2))),
"source_mae": float(np.mean(np.abs(output - observed))),
"percentile_mse": float(np.mean(rank_error**2)),
"percentile_mae": float(np.mean(np.abs(rank_error))),
"optimum_source_rmse": float(np.sqrt(np.mean((optimum - observed) ** 2))),
"optimum_source_mae": float(np.mean(np.abs(optimum - observed))),
}
metrics["noisy_targets"][key] = result
checks[f"{key}: predictions are finite"] = bool(np.isfinite(output).all())
checks[f"{key}: relative optimum error at most 0.25"] = relative <= 0.25
checks["Percentile and source-unit MSE population optima differ materially"] = bool(
np.mean(np.abs(targets["number_mse"] - targets["quantile_mse"]) / targets["number_mse"]) > 0.10
)
number_mse = np.asarray(metrics["noisy_targets"]["number_mse"]["group_predictions"])
quantile_mse = np.asarray(metrics["noisy_targets"]["quantile_mse"]["group_predictions"])
separation = float(np.mean((number_mse - quantile_mse) / targets["number_mse"]))
metrics["learned_mse_separation_fraction"] = separation
checks["Learned MSE predictions differ by at least 10 percent of source means"] = separation >= 0.10
return metrics, checks
if __name__ == "__main__":
report(PROOF_ID, run, seed=7701)
