1"""Local stage-2 track for branch-length-aware rooted-tree encoders."""
 2import numpy as np
 3
 4META = {
 5    "name": "rooted_tree_branch_classification",
 6    "domain": "graph-nn",
 7    "description": "Classification of rooted trees with identical edge-length multisets but different child topology; inputs encode parent indices and positive branch lengths.",
 8}
 9
10# Six vertices: root 0 and five descendants. Both classes have five edges.
11PARENTS = {
12    0: np.array([-1, 0, 0, 0, 3, 3]),  # root has 3 children; node 3 has 2
13    1: np.array([-1, 0, 1, 1, 1, 1]),  # root has 1 child; node 1 has 4
14}
15
16def _make(rng, n):
17    xs, ys = [], []
18    base = np.array([0.45, 0.72, 1.05, 1.38, 1.76], dtype=np.float32)
19    for _ in range(int(n)):
20        label = int(rng.randint(0, 2))
21        vals = rng.permutation(base)
22        # Small common nuisance scale preserves equal multiset within each sample.
23        vals = vals * rng.uniform(0.92, 1.08)
24        parent = PARENTS[label]
25        x = np.zeros((6, 2), dtype=np.float32)
26        x[:, 0] = np.where(parent < 0, -1.0, parent / 5.0)
27        edge_nodes = np.flatnonzero(parent >= 0)
28        x[edge_nodes, 1] = vals
29        xs.append(x); ys.append(label)
30    return np.asarray(xs), np.asarray(ys, dtype=np.int64)
31
32def get_dataset(seed, n_train, n_test):
33    xtr, ytr = _make(np.random.RandomState(int(seed)), n_train)
34    xte, yte = _make(np.random.RandomState(int(seed) + 100003), n_test)
35    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
36            "task": "classification", "metric": "err", "out_dim": 2}