1"""Custom structurally matched unordered point-set denoising track."""
 2import numpy as np
 3
 4META = {
 5    "name": "unordered_pointset_denoising",
 6    "domain": "point-set-diffusion",
 7    "description": "Denoising noisy unordered one-dimensional point configurations; sets are sorted only for canonical representation."
 8}
 9
10
11def _split(seed, n):
12    rng = np.random.RandomState(seed)
13    # Repulsive configurations: sorted Gaussian locations, with a mild global
14    # scale variation. The evaluation target is the clean set itself.
15    z = np.sort(rng.normal(size=(n, 6)).astype(np.float32), axis=1)
16    z = z / (np.std(z, axis=1, keepdims=True) + 0.35)
17    noisy = z + rng.normal(0, 0.22, size=z.shape).astype(np.float32)
18    noisy.sort(axis=1)
19    return noisy.astype(np.float32), z.astype(np.float32)
20
21
22def get_dataset(seed, n_train=400, n_test=160):
23    xtr, ytr = _split(seed, n_train)
24    xte, yte = _split(seed + 5000, n_test)
25    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
26            "task": "regression", "metric": "mse", "out_dim": 6}