1import json, math
 2import numpy as np
 3from block_tt3d import tt_svd_matrix, dense_from_cores, tt_apply, params
 4
 5
 6def relerr(a, b): return float(np.linalg.norm(a-b) / np.linalg.norm(a))
 7
 8
 9def main():
10    rng = np.random.default_rng(123)
11    modes = [2, 4, 4, 4]  # channel,x,y,z
12    N = int(np.prod(modes))
13    # Full-rank controlled-spectrum operator: no accidental exact low-rank result.
14    q, _ = np.linalg.qr(rng.normal(size=(N, N)))
15    p, _ = np.linalg.qr(rng.normal(size=(N, N)))
16    singular = np.geomspace(1.0, 1e-4, N)
17    A = q @ np.diag(singular) @ p.T
18    normA = np.linalg.norm(A)
19
20    tolerance = []
21    for frac in [0.50, 0.20, 0.10, 0.05, 0.01]:
22        eps = frac * normA
23        cores, ranks = tt_svd_matrix(A, modes, modes, eps=eps)
24        Ar = dense_from_cores(cores)
25        e = relerr(A, Ar)
26        tolerance.append({'relative_eps': frac, 'observed_relative_error': e,
27                          'bound_ratio': e/frac, 'ranks': ranks,
28                          'parameters': params(cores)})
29
30    rank_sweep = []
31    for r in [1, 2, 4, 8, 16, 32]:
32        cores, ranks = tt_svd_matrix(A, modes, modes, max_rank=r)
33        e = relerr(A, dense_from_cores(cores))
34        rank_sweep.append({'requested_rank': r, 'ranks': ranks,
35                           'parameters': params(cores), 'relative_error': e,
36                           'parameter_ratio_to_r2': params(cores)/(r*r)})
37
38    # Block-TT test: four channel blocks, each spatial 4x4x4 operator.
39    # Keep the same global matrix and compare independent block storage to one TT.
40    T = A.reshape(modes + modes)
41    block_parameters = 0
42    block_sqerr = 0.0
43    block_ranks = []
44    for o in range(2):
45        for i in range(2):
46            B = T[o, :, :, :, i, :, :, :].reshape(64, 64)
47            bc, br = tt_svd_matrix(B, modes[1:], modes[1:], eps=0.05*np.linalg.norm(B))
48            block_parameters += params(bc)
49            block_sqerr += np.linalg.norm(B-dense_from_cores(bc))**2
50            block_ranks.append(br)
51    mono, mr = tt_svd_matrix(A, modes, modes, eps=0.05*normA)
52    mono_rec = dense_from_cores(mono)
53
54    # Apply correctness and a tiny standard dense baseline memory comparison.
55    x = rng.normal(size=tuple(modes))
56    apply_error = relerr(A @ x.reshape(-1), tt_apply(mono, x).reshape(-1))
57    dense_parameters = A.size
58    result = {
59      'operator': {'shape': [N, N], 'dense_parameters': dense_parameters, 'frobenius_norm': float(normA)},
60      'prediction_checks': {
61        'svd_bound': 'relative error <= requested relative epsilon',
62        'rank_scaling': 'interior TT storage scales approximately as r^2',
63        'block_semantics': 'separate blocks can be rounded independently, but may cost more storage'},
64      'tolerance_sweep': tolerance,
65      'rank_sweep': rank_sweep,
66      'block_vs_monolithic': {
67        'block_parameters': block_parameters,
68        'block_relative_error': math.sqrt(block_sqerr)/normA,
69        'block_ranks': block_ranks,
70        'monolithic_parameters': params(mono),
71        'monolithic_relative_error': relerr(A, mono_rec),
72        'dense_parameters': dense_parameters},
73      'apply_relative_error': apply_error
74    }
75    print(json.dumps(result, indent=2))
76
77if __name__ == '__main__': main()