1import json
  2import time
  3from dataclasses import dataclass
  4import numpy as np
  5
  6SEED = 115
  7
  8@dataclass
  9class Node:
 10    lo: float
 11    hi: float
 12    slopes: list
 13    left: object = None
 14    right: object = None
 15    depth: int = 0
 16
 17    @property
 18    def splitting(self):
 19        return self.left is not None and self.right is not None
 20
 21
 22def build_tree(slopes, lo=0.0, hi=1.0, depth=0, max_depth=32):
 23    """Build a binary M-adic interval tree, retaining unary chains as shared nodes."""
 24    vals = sorted(float(x) for x in slopes)
 25    node = Node(lo, hi, vals, depth=depth)
 26    if len(vals) <= 1 or depth >= max_depth:
 27        return node
 28    mid = (lo + hi) / 2.0
 29    left = [x for x in vals if x < mid]
 30    right = [x for x in vals if x >= mid]
 31    if not left or not right:
 32        # A unary chain has no genuine branch decision.
 33        node.left = build_tree(vals, lo, hi, depth + 1, max_depth)
 34    else:
 35        node.left = build_tree(left, lo, mid, depth + 1, max_depth)
 36        node.right = build_tree(right, mid, hi, depth + 1, max_depth)
 37    return node
 38
 39
 40def split_counts(node, count=0):
 41    count += int(node.splitting)
 42    if node.left is None:
 43        return [count]
 44    if node.right is None:
 45        return split_counts(node.left, count)
 46    return split_counts(node.left, count) + split_counts(node.right, count)
 47
 48
 49def leaves(node):
 50    if node.left is None:
 51        return [node]
 52    if node.right is None:
 53        return leaves(node.left)
 54    return leaves(node.left) + leaves(node.right)
 55
 56
 57def tree_summary(root):
 58    counts = split_counts(root)
 59    return {"leaves": len(leaves(root)), "path_split_counts": counts,
 60            "max_splits": max(counts), "min_splits": min(counts)}
 61
 62
 63def lacunarity_ratios(slopes, alpha=0.0):
 64    vals = np.asarray(slopes, dtype=float)
 65    d = np.abs(vals - alpha)
 66    return d[1:] / np.maximum(d[:-1], 1e-12)
 67
 68
 69def sample_line(field, y, x, slope, radius=2):
 70    """Nearest-neighbour finite line average for offsets -radius..radius."""
 71    h, w = field.shape
 72    vals = []
 73    for t in range(-radius, radius + 1):
 74        yy = int(np.clip(round(y + t * slope), 0, h - 1))
 75        xx = int(np.clip(round(x + t), 0, w - 1))
 76        vals.append(field[yy, xx])
 77    return float(np.mean(np.abs(vals)))
 78
 79
 80def directional_responses(field, slopes, radius=2):
 81    h, w = field.shape
 82    out = np.empty((h, w, len(slopes)), dtype=np.float32)
 83    for y in range(h):
 84        for x in range(w):
 85            for k, slope in enumerate(slopes):
 86                out[y, x, k] = sample_line(field, y, x, slope, radius)
 87    return out
 88
 89
 90def routed_response(responses, slopes, root, beam=2):
 91    """Tree route: select the best beam child/subtree by cheap representative scores.
 92    The returned value uses only beam leaves, while dense responses serve as oracle data.
 93    """
 94    slope_to_k = {float(s): k for k, s in enumerate(slopes)}
 95    selected = []
 96    def visit(node, candidates):
 97        if node.left is None:
 98            selected.append(node.slopes[0]); return
 99        if node.right is None:
100            visit(node.left, candidates); return
101        children = [node.left, node.right]
102        scores = []
103        for child in children:
104            ks = [slope_to_k[float(s)] for s in child.slopes]
105            scores.append(float(np.mean(responses[..., ks])))
106        order = np.argsort(scores)[::-1][:min(beam, 2)]
107        for j in order:
108            visit(children[int(j)], candidates)
109    # Routing is applied globally in this compact MVP; leaf outputs are spatial maps.
110    visit(root, None)
111    selected = sorted(set(selected))
112    ks = [slope_to_k[s] for s in selected]
113    return np.max(responses[..., ks], axis=2), len(ks), selected
114
115
116def sparse_from_proxy(target, proxy, beam):
117    """Per-query routing using a cheap proxy, then evaluate only selected target lines."""
118    h, w, k = target.shape
119    chosen = np.argpartition(proxy, -beam, axis=2)[:, :, -beam:]
120    out = np.take_along_axis(target, chosen, axis=2).max(axis=2)
121    return out, int(h*w*beam)
122
123
124def run():
125    rng = np.random.default_rng(SEED)
126    k = 32
127    # Unique M-adic directions. For a binary tree, 32 distinct leaves necessarily
128    # require at least ceil(log2(32))=5 splitting vertices on some route.
129    slopes = np.linspace(0.0, 1.0, k, endpoint=False).tolist()
130    root = build_tree(slopes)
131    summary = tree_summary(root)
132    lac_seq = [0.5**j for j in range(1, 10)]
133    lac = lacunarity_ratios(lac_seq)
134    math_check = {
135        "geometric_max_ratio": float(lac.max()),
136        "geometric_expected_ratio": 0.5,
137        "geometric_ratio_allclose": bool(np.allclose(lac, 0.5)),
138        "tree_split_bound": summary["max_splits"],
139        "tree_paths": summary["path_split_counts"],
140        "binary_lower_bound_for_32_leaves": 5,
141        "bound_is_valid": summary["max_splits"] >= 5,
142    }
143    assert np.allclose(lac, 0.5)
144    assert summary["max_splits"] >= 5
145    h = w = 24
146    yy, xx = np.mgrid[:h, :w]
147    field = (np.sin((xx + 0.35 * yy) / 2.2) +
148             0.7 * np.cos((xx - 0.8 * yy) / 3.1) +
149             0.15 * rng.standard_normal((h, w))).astype(np.float32)
150    # Oracle target and a cheaper-radius proxy. The proxy is a routing surrogate,
151    # while target responses represent the expensive line-attention computation.
152    t0 = time.perf_counter(); target = directional_responses(field, slopes, radius=2); target_time = time.perf_counter()-t0
153    t0 = time.perf_counter(); proxy = directional_responses(field, slopes, radius=1); proxy_time = time.perf_counter()-t0
154    dense = np.max(target, axis=2)
155    idea, idea_samples = sparse_from_proxy(target, proxy, beam=2)
156    rng2 = np.random.default_rng(SEED)
157    random_idx = rng2.integers(0, k, size=(h, w, 2))
158    random_sparse = np.take_along_axis(target, random_idx, axis=2).max(axis=2)
159    def rel_error(a):
160        return float(np.mean(np.abs(a-dense)) / (np.mean(np.abs(dense))+1e-8))
161    metrics = {
162        "K": k, "tree": summary,
163        "dense_samples": int(h*w*k), "idea_target_samples": idea_samples,
164        "random_target_samples": idea_samples,
165        "sample_reduction": float(k/2),
166        "dense_target_ms": 1000*target_time,
167        "proxy_routing_ms": 1000*proxy_time,
168        "idea_relative_error": rel_error(idea),
169        "random_relative_error": rel_error(random_sparse),
170        "idea_better_than_random": rel_error(idea) < rel_error(random_sparse),
171    }
172    result = {"seed": SEED, "math_check": math_check, "metrics": metrics}
173    with open("results.json", "w") as f:
174        json.dump(result, f, indent=2)
175    print(json.dumps(result, indent=2))
176
177if __name__ == "__main__":
178    run()