|
| 1 | +""" |
| 2 | +Demonstration of hyperedge support using cotengra in TensorCircuit. |
| 3 | +""" |
| 4 | + |
| 5 | +import time |
| 6 | +import numpy as np |
| 7 | +import tensornetwork as tn |
| 8 | +import tensorcircuit as tc |
| 9 | + |
| 10 | + |
| 11 | +def hyperedge_demo(): |
| 12 | + print("Demonstrating hyperedge contraction with cotengra...") |
| 13 | + |
| 14 | + # 1. Single Hyperedge Example |
| 15 | + # Three tensors A, B, C connected by a single hyperedge (CopyNode) |
| 16 | + # Result should be sum_i A_i * B_i * C_i |
| 17 | + |
| 18 | + dim = 2 |
| 19 | + a = tn.Node(np.array([1.0, 2.0]), name="A") |
| 20 | + b = tn.Node(np.array([1.0, 2.0]), name="B") |
| 21 | + c = tn.Node(np.array([1.0, 2.0]), name="C") |
| 22 | + cn = tn.CopyNode(3, dim, name="CN") |
| 23 | + |
| 24 | + a[0] ^ cn[0] |
| 25 | + b[0] ^ cn[1] |
| 26 | + c[0] ^ cn[2] |
| 27 | + |
| 28 | + nodes = [a, b, c, cn] |
| 29 | + |
| 30 | + # Set contractor to cotengra |
| 31 | + tc.set_contractor("cotengra") |
| 32 | + |
| 33 | + res = tc.contractor(nodes) |
| 34 | + print("Single Hyperedge Result:", res.tensor) |
| 35 | + expected = 1 * 1 * 1 + 2 * 2 * 2 |
| 36 | + print(f"Expected: {expected}") |
| 37 | + assert np.allclose(res.tensor, expected) |
| 38 | + |
| 39 | + # 2. Large Scale Hyperedge Example |
| 40 | + # Demonstrate memory and time efficiency with a large number of legs |
| 41 | + print("\nDemonstrating large scale hyperedge (20 legs)...") |
| 42 | + num_legs = 20 |
| 43 | + dim = 2 |
| 44 | + |
| 45 | + # Create 20 random tensors connected to a single CopyNode |
| 46 | + input_tensors = [ |
| 47 | + tn.Node(np.random.rand(dim), name=f"T{i}") for i in range(num_legs) |
| 48 | + ] |
| 49 | + cn_large = tn.CopyNode(num_legs, dim, name="CN_Large") |
| 50 | + |
| 51 | + for i, t in enumerate(input_tensors): |
| 52 | + t[0] ^ cn_large[i] |
| 53 | + |
| 54 | + large_nodes = input_tensors + [cn_large] |
| 55 | + |
| 56 | + start_time = time.time() |
| 57 | + res_large = tc.contractor(large_nodes) |
| 58 | + end_time = time.time() |
| 59 | + |
| 60 | + print(f"Contracted {num_legs} legs in {end_time - start_time:.4f} seconds.") |
| 61 | + print("Large Hyperedge Result shape:", res_large.tensor.shape) |
| 62 | + |
| 63 | + # Verification: Explicitly calculate the sum |
| 64 | + # result = sum_k (prod_i T_i[k]) |
| 65 | + |
| 66 | + # Transpose input tensors to shape (num_legs, dim) |
| 67 | + tensor_matrix = np.stack([t.tensor for t in input_tensors]) |
| 68 | + # Product along the tensor axis (0) for each dimension index |
| 69 | + prod_along_legs = np.prod(tensor_matrix, axis=0) |
| 70 | + expected_sum = np.sum(prod_along_legs) |
| 71 | + |
| 72 | + print(f"Computed: {res_large.tensor}") |
| 73 | + print(f"Expected: {expected_sum}") |
| 74 | + assert np.allclose(res_large.tensor, expected_sum) |
| 75 | + |
| 76 | + |
| 77 | +if __name__ == "__main__": |
| 78 | + hyperedge_demo() |
0 commit comments