import sys, json, math from pathlib import Path import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report, count_params # A small genuine Cartesian 3-D operator task: two physical channels on 4^3 voxels. META = {'name': 'voxel_operator_3d', 'domain': '3d_operator', 'description': 'Learn a channel-blocked linear operator on a 2x4x4x4 Cartesian voxel grid.'} def _tt_from_dense(a, modes, rank=2): # TT-SVD for a matrix whose output/input modes are both modes. d = len(modes) t = a.reshape(tuple(modes) + tuple(modes)) t = np.transpose(t, sum(([k, d+k] for k in range(d)), [])) work = t.reshape([modes[k] * modes[k] for k in range(d)]) cores, ranks = [], [1] for k in range(d - 1): work = work.reshape(ranks[-1] * modes[k] * modes[k], -1) u, s, vt = np.linalg.svd(work, full_matrices=False) r = min(rank, len(s)) cores.append((u[:, :r]).reshape(ranks[-1], modes[k], modes[k], r)) work = s[:r, None] * vt[:r] ranks.append(r) cores.append(work.reshape(ranks[-1], modes[-1], modes[-1], 1)) return cores def get_dataset(seed, n_train=400, n_test=100): rng = np.random.default_rng(seed) modes = (2, 4, 4, 4) # Block semantics: diagonal and cross-channel spatial maps have distinct TT structure. blocks = [] for q in range(4): b = np.zeros((64, 64)) for r in range(2): fs = [rng.normal(size=(4, 4)) for _ in range(3)] b += np.kron(np.kron(fs[0], fs[1]), fs[2]) blocks.append(b * (0.65 if q < 2 else 0.28)) target = np.zeros((128, 128)) for o in range(2): for i in range(2): target[o*64:(o+1)*64, i*64:(i+1)*64] = blocks[2*o+i] def sample(n, s): z = np.random.default_rng(s).normal(size=(n, 128)).astype('float32') y = (z @ target.T).sum(axis=1, keepdims=True) / 12.0 y += 0.01 * np.random.default_rng(s + 99).normal(size=y.shape) return z.reshape(n, *modes).astype('float32'), y.astype('float32') xtr, ytr = sample(n_train, seed + 10) xte, yte = sample(n_test, seed + 5010) return {'xtr': torch.from_numpy(xtr), 'ytr': torch.from_numpy(ytr), 'xte': torch.from_numpy(xte), 'yte': torch.from_numpy(yte), 'task': 'regression', 'metric': 'mse', 'input_shape': modes, 'out_dim': 1, 'track': META['name']} class DenseOperator(nn.Module): def __init__(self): super().__init__() self.op = nn.Linear(128, 128, bias=False) self.head = nn.Linear(128, 1) def forward(self, x): return self.head(self.op(x.flatten(1))) class TTBlockOperator(nn.Module): def __init__(self, rank=2): super().__init__() self.modes = (4, 4, 4) self.rank = rank # Four independent channel-pair TT matrices: explicit semantic blocks. self.cores = nn.ParameterList() for _ in range(4): rs = [1, rank, rank, 1] for k, n in enumerate(self.modes): p = nn.Parameter(torch.randn(rs[k], n, n, rs[k+1]) * (0.12 if k else 0.25)) self.cores.append(p) self.head = nn.Linear(128, 1) def block_apply(self, x, offset): # x: [B,4,4,4], output same spatial grid; no dense matrix materialization. gs = [self.cores[offset*3+k] for k in range(3)] return torch.einsum('tabc,piaq,qjbr,rkcs->tijk', x, gs[0], gs[1], gs[2]).reshape(x.shape[0], 64) def forward(self, x): x = x.reshape(x.shape[0], 2, 64) ys = [] for o in range(2): y = 0.0 for i in range(2): y = y + self.block_apply(x[:, i].reshape(-1,4,4,4), 2*o+i) ys.append(y) return self.head(torch.cat(ys, dim=1)) def make_train(kind, lr, rank=2): def run(seed): torch.manual_seed(seed + 123) d = get_dataset(seed) model = DenseOperator() if kind == 'dense' else TTBlockOperator(rank) _, metric, _ = train_model(model, d, epochs=20, lr=lr, batch=128, weight_decay=0.0, log=lambda *_: None) return metric return run def signature(): # Signature is measured on a model actually trained on the benchmark task. torch.manual_seed(7) d = get_dataset(0) m = TTBlockOperator(2) m, _, _ = train_model(m, d, epochs=20, lr=1e-2, batch=128, weight_decay=0.0, log=lambda *_: None) m.eval() device = next(m.parameters()).device x = d['xte'][:16].to(device) with torch.no_grad(): y = m(x) full = torch.zeros(128, 128, dtype=x.dtype, device=device) eye = torch.eye(64, device=device, dtype=x.dtype).reshape(64, 4, 4, 4) for o in range(2): for i in range(2): cols = [m.block_apply(eye[q:q+1], 2*o+i).squeeze(0) for q in range(64)] block = torch.stack(cols, dim=1) full[o*64:(o+1)*64, i*64:(i+1)*64] = block explicit = m.head(x.reshape(16, 128) @ full.T) rel = float((y-explicit).norm() / (y.norm()+1e-12)) dense_operator_params = 128*128 tt_operator_params = sum(p.numel() for p in m.cores) return {'prediction': 'trained block TT preserves its sequential operator map with reduced operator storage', 'predicted_operator_params': dense_operator_params, 'observed_trained_tt_operator_params': tt_operator_params, 'observed_storage_ratio': tt_operator_params / float(dense_operator_params), 'observed_trained_forward_relative_error': rel, 'confirmed': bool(tt_operator_params < dense_operator_params and rel < 1e-6)} def main(): # Union parity: every idea lr is also evaluated by the baseline. grid = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}] base = sweep_baseline(lambda c: make_train('dense', c['lr']), grid, seeds=(0,1,2,3)) # Equal-budget idea sweep: baseline's best lr plus two nearby values; all are # already present in the baseline grid (search-space parity). idea_runs = [] for cfg in grid: r = evaluate(make_train('tt', cfg['lr']), seeds=range(8)) idea_runs.append({'cfg': cfg, 'result': r}) best = min(idea_runs, key=lambda z: z['result']['mean']) idea = dict(best['result']) idea['best_cfg'] = best['cfg'] idea['configs_tried'] = [{'cfg': z['cfg'], 'mean': z['result']['mean']} for z in idea_runs] rep = make_report(META['name'], 'dense_vs_block_tt', base, idea, {'mechanism_signature': signature(), 'custom_track': {'name': META['name'], 'file': 'voxel_operator_track.py', 'domain': '3d_operator'}, 'protocol_note': '8 paired seeds; custom 3-D Cartesian track; baseline sweep and shared lr union.'}) print(json.dumps(rep, indent=2)) if __name__ == '__main__': main()