1import json, math
  2import numpy as np
  3
  4
  5def tt_svd_matrix(A, out_shape, in_shape, eps=0.0, max_rank=None):
  6    d = len(out_shape)
  7    T = np.asarray(A).reshape(tuple(out_shape) + tuple(in_shape))
  8    # Interleave (output,input) so each TT core owns one Cartesian mode pair.
  9    T = np.transpose(T, sum(([k, d+k] for k in range(d)), []))
 10    T = T.reshape([out_shape[k]*in_shape[k] for k in range(d)])
 11    cores, ranks = [], [1]
 12    work = T
 13    total_sq = np.linalg.norm(work.ravel())**2
 14    for k in range(d-1):
 15        work = work.reshape(ranks[-1]*out_shape[k]*in_shape[k], -1)
 16        U, s, Vt = np.linalg.svd(work, full_matrices=False)
 17        # Distribute a global tolerance over d-1 cuts, as in the TT error bound.
 18        budget = (eps / math.sqrt(max(1, d-1)))**2
 19        tail = np.cumsum(s[::-1]**2)[::-1]
 20        keep = len(s) if eps == 0 else int(np.searchsorted(tail <= budget, True))
 21        # searchsorted on a boolean is awkward for nonmonotonic bool; explicit minimum.
 22        if eps != 0:
 23            valid = np.where(tail <= budget)[0]
 24            keep = int(valid[0]+1) if len(valid) else len(s)
 25        if max_rank is not None: keep = min(keep, max_rank)
 26        keep = max(1, keep)
 27        cores.append(U[:, :keep].reshape(ranks[-1], out_shape[k], in_shape[k], keep))
 28        work = (s[:keep,None] * Vt[:keep])
 29        ranks.append(keep)
 30    cores.append(work.reshape(ranks[-1], out_shape[-1], in_shape[-1], 1))
 31    return cores, ranks
 32
 33
 34def tt_reconstruct(cores):
 35    x = cores[0]
 36    for g in cores[1:]:
 37        x = np.tensordot(x, g, axes=([-1],[0]))
 38    # x: (1,o1,i1,...,od,id,1), remove boundaries then regroup
 39    x = np.squeeze(x, (0,-1))
 40    d = len(cores)
 41    x = np.transpose(x, list(range(0, 2*d, 2)) + list(range(1, 2*d, 2)))
 42    shape = [cores[k].shape[1] for k in range(d)] + [cores[k].shape[2] for k in range(d)]
 43    return x.reshape(shape)
 44
 45
 46def tt_apply(cores, x):
 47    """Apply a TT matrix by contracting input modes and neighboring TT ranks."""
 48    d = len(cores)
 49    # Unique einsum symbols: input/output modes plus TT ranks.
 50    alphabet = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
 51    need = 3*d + 1
 52    if need > len(alphabet): raise ValueError('too many modes')
 53    inp = alphabet[:d]
 54    out = alphabet[d:2*d]
 55    ranks = alphabet[2*d:3*d+1]
 56    terms = [inp]
 57    for k, g in enumerate(cores):
 58        terms.append(ranks[k] + out[k] + inp[k] + ranks[k+1])
 59    equation = ','.join(terms) + '->' + out
 60    return np.einsum(equation, x, *cores, optimize=True)
 61
 62def dense_from_cores(cores):
 63    t = tt_reconstruct(cores)
 64    d = len(cores); os=t.shape[:d]; ins=t.shape[d:]
 65    return t.reshape(int(np.prod(os)), int(np.prod(ins)))
 66
 67
 68def params(cores): return int(sum(g.size for g in cores))
 69
 70
 71def make_spatial_operator(seed=4, n=4, rank=2):
 72    rng=np.random.default_rng(seed)
 73    # CP sum of separable 1D factors; exact spatial TT rank <= rank.
 74    S=np.zeros((n,n,n,n,n,n))
 75    for q in range(rank):
 76        fs=[rng.normal(size=(n,n)) for _ in range(3)]
 77        term=fs[0][:,:,None,None,None,None]*fs[1][None,None,:,:,None,None]*fs[2][None,None,None,None,:,:]
 78        S += term
 79    return S
 80
 81
 82def main():
 83    rng=np.random.default_rng(7)
 84    n=4; C=2; spatial=make_spatial_operator(n=n,rank=2)
 85    # Channel semantics: diagonal blocks share one spatial map, off-diagonal use a distinct map.
 86    spatial2=make_spatial_operator(seed=12,n=n,rank=2)
 87    # tensor mode order (channel,x,y,z), matrix axes output then input
 88    A=np.zeros((C,n,n,n,C,n,n,n))
 89    A[0, :, :, :, 0] = spatial
 90    A[1, :, :, :, 1] = spatial
 91    A[0, :, :, :, 1] = 0.35*spatial2
 92    A[1, :, :, :, 0] = -0.25*spatial2
 93    dense=A.reshape(C*n**3,C*n**3)
 94
 95    # Prediction 1: TT-SVD tolerance gives global Frobenius error <= eps (up to roundoff).
 96    tol_rows=[]
 97    for eps in [1e-1, 3e-2, 1e-2, 3e-3]:
 98        cs,rs=tt_svd_matrix(dense,[C,n,n,n],[C,n,n,n],eps=eps)
 99        rec=dense_from_cores(cs)
100        rel=np.linalg.norm(dense-rec)/np.linalg.norm(dense)
101        tol_rows.append({'eps':eps,'relative_error':float(rel),'ranks':rs,'bound_ratio':float(rel/eps)})
102
103    # Prediction 2: uniform TT parameter count is affine/quadratic in intermediate rank.
104    rank_rows=[]
105    for r in [1,2,3,4,6]:
106        cs,rs=tt_svd_matrix(dense,[C,n,n,n],[C,n,n,n],max_rank=r)
107        rank_rows.append({'requested_rank':r,'actual_ranks':rs,'parameters':params(cs),'param_over_r2':params(cs)/(r*r)})
108
109    # Prediction 3: block representation preserves independent channel blocks.
110    # Decompose each of four channel blocks as a spatial TT and sum storage.
111    block_p=0; block_err=0.; block_details=[]
112    for o in range(C):
113        for i in range(C):
114            b=A[o,:,:,:,i,:,:,:].reshape(n**3,n**3)
115            bc,br=tt_svd_matrix(b,[n,n,n],[n,n,n],eps=1e-10)
116            block_p += params(bc)
117            block_err += np.linalg.norm(b-dense_from_cores(bc))**2
118            block_details.append(br)
119    mono, mr=tt_svd_matrix(dense,[C,n,n,n],[C,n,n,n],eps=1e-10)
120    mono_err=np.linalg.norm(dense-dense_from_cores(mono))/np.linalg.norm(dense)
121
122    # Direct contraction correctness on a small random tensor.
123    x=rng.normal(size=(C,n,n,n)); y1=dense@x.reshape(-1); y2=tt_apply(mono,x).reshape(-1)
124    apply_rel=np.linalg.norm(y1-y2)/np.linalg.norm(y1)
125    result={'tolerance_sweep':tol_rows,'rank_sweep':rank_rows,
126            'block_vs_monolithic':{'block_parameters':block_p,'block_relative_error':float(math.sqrt(block_err)/np.linalg.norm(dense)),
127              'block_ranks':block_details,'monolithic_parameters':params(mono),'monolithic_ranks':mr,'monolithic_relative_error':float(mono_err)},
128            'apply_relative_error':float(apply_rel),
129            'predictions':{'tolerance':'relative error should be <= eps','rank':'parameters should grow approximately quadratically in interior rank','blocks':'block TT should represent independent channel blocks without channel rank mixing'}}
130    print(json.dumps(result,indent=2))
131
132if __name__=='__main__': main()