1import sys, json, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report, count_params
  9
 10# A small genuine Cartesian 3-D operator task: two physical channels on 4^3 voxels.
 11META = {'name': 'voxel_operator_3d', 'domain': '3d_operator',
 12        'description': 'Learn a channel-blocked linear operator on a 2x4x4x4 Cartesian voxel grid.'}
 13
 14
 15def _tt_from_dense(a, modes, rank=2):
 16    # TT-SVD for a matrix whose output/input modes are both modes.
 17    d = len(modes)
 18    t = a.reshape(tuple(modes) + tuple(modes))
 19    t = np.transpose(t, sum(([k, d+k] for k in range(d)), []))
 20    work = t.reshape([modes[k] * modes[k] for k in range(d)])
 21    cores, ranks = [], [1]
 22    for k in range(d - 1):
 23        work = work.reshape(ranks[-1] * modes[k] * modes[k], -1)
 24        u, s, vt = np.linalg.svd(work, full_matrices=False)
 25        r = min(rank, len(s))
 26        cores.append((u[:, :r]).reshape(ranks[-1], modes[k], modes[k], r))
 27        work = s[:r, None] * vt[:r]
 28        ranks.append(r)
 29    cores.append(work.reshape(ranks[-1], modes[-1], modes[-1], 1))
 30    return cores
 31
 32
 33def get_dataset(seed, n_train=400, n_test=100):
 34    rng = np.random.default_rng(seed)
 35    modes = (2, 4, 4, 4)
 36    # Block semantics: diagonal and cross-channel spatial maps have distinct TT structure.
 37    blocks = []
 38    for q in range(4):
 39        b = np.zeros((64, 64))
 40        for r in range(2):
 41            fs = [rng.normal(size=(4, 4)) for _ in range(3)]
 42            b += np.kron(np.kron(fs[0], fs[1]), fs[2])
 43        blocks.append(b * (0.65 if q < 2 else 0.28))
 44    target = np.zeros((128, 128))
 45    for o in range(2):
 46        for i in range(2):
 47            target[o*64:(o+1)*64, i*64:(i+1)*64] = blocks[2*o+i]
 48    def sample(n, s):
 49        z = np.random.default_rng(s).normal(size=(n, 128)).astype('float32')
 50        y = (z @ target.T).sum(axis=1, keepdims=True) / 12.0
 51        y += 0.01 * np.random.default_rng(s + 99).normal(size=y.shape)
 52        return z.reshape(n, *modes).astype('float32'), y.astype('float32')
 53    xtr, ytr = sample(n_train, seed + 10)
 54    xte, yte = sample(n_test, seed + 5010)
 55    return {'xtr': torch.from_numpy(xtr), 'ytr': torch.from_numpy(ytr),
 56            'xte': torch.from_numpy(xte), 'yte': torch.from_numpy(yte),
 57            'task': 'regression', 'metric': 'mse', 'input_shape': modes,
 58            'out_dim': 1, 'track': META['name']}
 59
 60
 61class DenseOperator(nn.Module):
 62    def __init__(self):
 63        super().__init__()
 64        self.op = nn.Linear(128, 128, bias=False)
 65        self.head = nn.Linear(128, 1)
 66    def forward(self, x):
 67        return self.head(self.op(x.flatten(1)))
 68
 69
 70class TTBlockOperator(nn.Module):
 71    def __init__(self, rank=2):
 72        super().__init__()
 73        self.modes = (4, 4, 4)
 74        self.rank = rank
 75        # Four independent channel-pair TT matrices: explicit semantic blocks.
 76        self.cores = nn.ParameterList()
 77        for _ in range(4):
 78            rs = [1, rank, rank, 1]
 79            for k, n in enumerate(self.modes):
 80                p = nn.Parameter(torch.randn(rs[k], n, n, rs[k+1]) * (0.12 if k else 0.25))
 81                self.cores.append(p)
 82        self.head = nn.Linear(128, 1)
 83    def block_apply(self, x, offset):
 84        # x: [B,4,4,4], output same spatial grid; no dense matrix materialization.
 85        gs = [self.cores[offset*3+k] for k in range(3)]
 86        return torch.einsum('tabc,piaq,qjbr,rkcs->tijk', x, gs[0], gs[1], gs[2]).reshape(x.shape[0], 64)
 87    def forward(self, x):
 88        x = x.reshape(x.shape[0], 2, 64)
 89        ys = []
 90        for o in range(2):
 91            y = 0.0
 92            for i in range(2):
 93                y = y + self.block_apply(x[:, i].reshape(-1,4,4,4), 2*o+i)
 94            ys.append(y)
 95        return self.head(torch.cat(ys, dim=1))
 96
 97
 98def make_train(kind, lr, rank=2):
 99    def run(seed):
100        torch.manual_seed(seed + 123)
101        d = get_dataset(seed)
102        model = DenseOperator() if kind == 'dense' else TTBlockOperator(rank)
103        _, metric, _ = train_model(model, d, epochs=20, lr=lr, batch=128, weight_decay=0.0, log=lambda *_: None)
104        return metric
105    return run
106
107
108def signature():
109    # Signature is measured on a model actually trained on the benchmark task.
110    torch.manual_seed(7)
111    d = get_dataset(0)
112    m = TTBlockOperator(2)
113    m, _, _ = train_model(m, d, epochs=20, lr=1e-2, batch=128,
114                           weight_decay=0.0, log=lambda *_: None)
115    m.eval()
116    device = next(m.parameters()).device
117    x = d['xte'][:16].to(device)
118    with torch.no_grad():
119        y = m(x)
120        full = torch.zeros(128, 128, dtype=x.dtype, device=device)
121        eye = torch.eye(64, device=device, dtype=x.dtype).reshape(64, 4, 4, 4)
122        for o in range(2):
123            for i in range(2):
124                cols = [m.block_apply(eye[q:q+1], 2*o+i).squeeze(0)
125                        for q in range(64)]
126                block = torch.stack(cols, dim=1)
127                full[o*64:(o+1)*64, i*64:(i+1)*64] = block
128        explicit = m.head(x.reshape(16, 128) @ full.T)
129        rel = float((y-explicit).norm() / (y.norm()+1e-12))
130        dense_operator_params = 128*128
131        tt_operator_params = sum(p.numel() for p in m.cores)
132    return {'prediction': 'trained block TT preserves its sequential operator map with reduced operator storage',
133            'predicted_operator_params': dense_operator_params,
134            'observed_trained_tt_operator_params': tt_operator_params,
135            'observed_storage_ratio': tt_operator_params / float(dense_operator_params),
136            'observed_trained_forward_relative_error': rel,
137            'confirmed': bool(tt_operator_params < dense_operator_params and rel < 1e-6)}
138
139
140def main():
141    # Union parity: every idea lr is also evaluated by the baseline.
142    grid = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}]
143    base = sweep_baseline(lambda c: make_train('dense', c['lr']), grid,
144                          seeds=(0,1,2,3))
145    # Equal-budget idea sweep: baseline's best lr plus two nearby values; all are
146    # already present in the baseline grid (search-space parity).
147    idea_runs = []
148    for cfg in grid:
149        r = evaluate(make_train('tt', cfg['lr']), seeds=range(8))
150        idea_runs.append({'cfg': cfg, 'result': r})
151    best = min(idea_runs, key=lambda z: z['result']['mean'])
152    idea = dict(best['result'])
153    idea['best_cfg'] = best['cfg']
154    idea['configs_tried'] = [{'cfg': z['cfg'], 'mean': z['result']['mean']} for z in idea_runs]
155    rep = make_report(META['name'], 'dense_vs_block_tt', base, idea,
156                      {'mechanism_signature': signature(),
157                       'custom_track': {'name': META['name'], 'file': 'voxel_operator_track.py', 'domain': '3d_operator'},
158                       'protocol_note': '8 paired seeds; custom 3-D Cartesian track; baseline sweep and shared lr union.'})
159    print(json.dumps(rep, indent=2))
160
161if __name__ == '__main__':
162    main()