add hypernode support for contractor

Author: refraction-ray

examples/contraction_policy_benchmark.py

@@ -0,0 +1,98 @@
+"""
+Benchmark JAX JIT staging and execution time for VQE for two different contraction policies.
+Comparing use_primitives=True (Algebraic Path) vs use_primitives=False (Legacy Path).
+"""
+
+import time
+import jax
+import jax.numpy as jnp
+import tensorcircuit as tc
+
+# Set backend to JAX
+tc.set_backend("jax")
+
+
+def run_vqe_benchmark(n, nlayers, use_primitives):
+    # Set the contractor configuration
+    tc.set_contractor("cotengra", use_primitives=use_primitives)
+
+    def energy_fn(params):
+        c = tc.Circuit(n)
+        idx = 0
+        for _ in range(nlayers):
+            # Single qubit rotations
+            for i in range(n):
+                c.rx(i, theta=params[idx])
+                idx += 1
+                c.ry(i, theta=params[idx])
+                idx += 1
+            # Entangling Rzz gates
+            for i in range(n - 1):
+                c.rzz(i, i + 1, theta=params[idx])
+                idx += 1
+
+        # Compute expectation value of a TFIM Hamiltonian
+        e = 0.0
+        for i in range(n - 1):
+            e += c.expectation_ps(z=[i, i + 1])
+        for i in range(n):
+            e += c.expectation_ps(x=[i])
+        return jnp.real(e)
+
+    # Use value_and_grad for a more realistic VQE benchmark
+    val_grad_fn = jax.jit(jax.value_and_grad(energy_fn))
+
+    # Prepare random parameters
+    num_params = nlayers * (2 * n + (n - 1))
+    params = jax.random.normal(jax.random.PRNGKey(42), (num_params,))
+
+    # 1. Staging Time (Compilation + First Execution)
+    start = time.time()
+    v, g = val_grad_fn(params)
+    v.block_until_ready()
+    # Gradient nodes are usually ready when the value is ready if fused,
+    # but we can block on jnp.sum(g) to be sure
+    jnp.sum(g).block_until_ready()
+    staging_time = time.time() - start
+
+    # 2. Execution Time (Subsequent Runs)
+    iters = 10
+    start = time.time()
+    for _ in range(iters):
+        v, g = val_grad_fn(params)
+        v.block_until_ready()
+        jnp.sum(g).block_until_ready()
+    exec_time = (time.time() - start) / iters
+
+    return staging_time, exec_time
+
+
+def main():
+    n = 12
+    nlayers = 5
+    print(f"--- VQE Benchmark: {n} Qubits, {nlayers} Layers ---")
+    print(f"Backend: {tc.backend.name}")
+
+    # Test use_primitives=False (Legacy Path)
+    print("\n[Case 1] use_primitives=False (Legacy Path)")
+    s1, e1 = run_vqe_benchmark(n, nlayers, use_primitives=False)
+    print(f"  Staging Time:   {s1:.4f}s")
+    print(f"  Execution Time: {e1:.6f}s")
+
+    # Test use_primitives=True (Algebraic Path)
+    print("\n[Case 2] use_primitives=True (Algebraic Path)")
+    s2, e2 = run_vqe_benchmark(n, nlayers, use_primitives=True)
+    print(f"  Staging Time:   {s2:.4f}s")
+    print(f"  Execution Time: {e2:.6f}s")
+
+    print("\n--- Summary ---")
+    staging_gain = (s1 - s2) / s1 * 100 if s1 > 0 else 0
+    exec_gain = (e1 - e2) / e1 * 100 if e1 > 0 else 0
+    print(f"Staging Speedup:   {staging_gain:.2f}%")
+    print(f"Execution Speedup: {exec_gain:.2f}%")
+
+
+if __name__ == "__main__":
+    main()
+
+# to me, the diff is not significant

examples/hyperedge_demo.py

@@ -0,0 +1,143 @@
+"""
+Physics-relevant demonstration of hyperedge support in TensorCircuit.
+Computing the partition function of a 2D classical Ising model using CopyNodes.
+"""
+
+import time
+import jax
+import jax.numpy as jnp
+import tensornetwork as tn
+import tensorcircuit as tc
+
+# Set backend to JAX for JIT and AD support
+tc.set_backend("jax")
+tc.set_dtype("complex128")
+tc.set_contractor("cotengra")
+
+
+def ising_partition_function(L, beta, J=1.0):
+    """
+    Compute the partition function of a 2D Ising model on an L x L grid.
+    Uses CopyNodes to represent spins and 2-index tensors for Boltzmann factors.
+
+    The partition function is Z = sum_{s} exp(beta * J * sum_{<i,j>} s_i * s_j).
+    Each site i has a spin s_i in {1, -1}.
+    Each bond <i,j> contributes a factor exp(beta * J * s_i * s_j).
+    """
+    # Boltzmann factor matrix M_{si, sj} = exp(beta * J * si * sj)
+    # spins are {1, -1}, mapped to indices {0, 1}
+    # si*sj = 1 if si==sj (indices 00 or 11), -1 if si!=sj (indices 01 or 10)
+    # Using jnp.exp to allow AD through beta
+    M = jnp.array(
+        [
+            [jnp.exp(beta * J), jnp.exp(-beta * J)],
+            [jnp.exp(-beta * J), jnp.exp(beta * J)],
+        ]
+    )
+
+    nodes = []
+    # Grid of CopyNodes (Delta tensors) representing the spins
+    grid = [[None for _ in range(L)] for _ in range(L)]
+
+    for i in range(L):
+        for j in range(L):
+            # Determine degree of CopyNode based on neighbors (open BC)
+            degree = 0
+            if i > 0:
+                degree += 1
+            if i < L - 1:
+                degree += 1
+            if j > 0:
+                degree += 1
+            if j < L - 1:
+                degree += 1
+
+            # A CopyNode(degree, 2) enforces that all connected legs have the same value (spin state)
+            cn = tn.CopyNode(degree, 2, name=f"site_{i}_{j}")
+            grid[i][j] = cn
+
+    # Track which axis of each CopyNode is used as we connect bonds
+    axis_ptr = [[0 for _ in range(L)] for _ in range(L)]
+
+    # Add bond tensors and connect to the CopyNodes
+    for i in range(L):
+        for j in range(L):
+            # Horizontal bond to the right
+            if j < L - 1:
+                bond_h = tn.Node(M, name=f"bond_h_{i}_{j}")
+                nodes.append(bond_h)
+                grid[i][j][axis_ptr[i][j]] ^ bond_h[0]
+                grid[i][j + 1][axis_ptr[i][j + 1]] ^ bond_h[1]
+                axis_ptr[i][j] += 1
+                axis_ptr[i][j + 1] += 1
+
+            # Vertical bond downwards
+            if i < L - 1:
+                bond_v = tn.Node(M, name=f"bond_v_{i}_{j}")
+                nodes.append(bond_v)
+                grid[i][j][axis_ptr[i][j]] ^ bond_v[0]
+                grid[i + 1][j][axis_ptr[i + 1][j]] ^ bond_v[1]
+                axis_ptr[i][j] += 1
+                axis_ptr[i + 1][j] += 1
+
+    # Multi-node contraction with cotengra (which handles hyperedges efficiently)
+    # The algebraic path is triggered automatically because CopyNodes are present.
+    all_nodes = nodes + [grid[i][j] for i in range(L) for j in range(L)]
+
+    # Ensure cotengra is used for high-performance contraction
+    z_node = tc.contractor(all_nodes)
+
+    return z_node.tensor
+
+
+def main():
+    L = 8
+    J = 1.0
+    beta = 0.4  # Near critical point beta_c approx 0.44 for 2D Ising
+    print(f"--- 2D Ising Model Partition Function ({L}x{L} lattice) ---")
+    print(f"Backend: {tc.backend.name}, Parameters: J={J}, beta={beta}")
+
+    # 1. Direct computation
+    z = ising_partition_function(L, beta, J)
+    print(f"Z({beta}) = {z:.6f}")
+
+    # 2. JIT-compiled version
+    # JIT significantly accelerates repeated calls with the same topology
+    print("\nDemonstrating JIT acceleration...")
+    ising_jit = jax.jit(ising_partition_function, static_argnums=(0,))
+
+    start = time.time()
+    _ = ising_jit(L, beta, J)
+    print(f"First run (with JIT warmup): {time.time() - start:.4f}s")
+
+    start = time.time()
+    _ = ising_jit(L, beta, J)
+    print(f"Second run (JIT cached):     {time.time() - start:.4f}s")
+
+    # 3. Automatic Differentiation (AD)
+    # Internal Energy U = - d(ln Z) / d(beta)
+    print("\nComputing Internal Energy via Automatic Differentiation...")
+
+    def log_z(beta_val):
+        # We take the real part as the partition function is real
+        val = ising_partition_function(L, beta_val, J)
+        return jnp.log(tc.backend.real(val))
+
+    energy_fn = jax.grad(log_z)
+    energy = energy_fn(beta)
+
+    # U = -d(ln Z)/d(beta) in our convention
+    print(f"Expectation of Internal Energy <E> = {-energy:.6f}")
+
+    # 4. Scaling demonstration
+    L_larger = 12
+    print(f"\nScaling check: L={L_larger} ({L_larger*L_larger} spins)...")
+    start = time.time()
+    z_large = ising_jit(L_larger, beta, J)
+    print(
+        f"Result for L={L_larger}: {z_large:.2e} (computed in {time.time()-start:.4f}s)"
+    )
+
+
+if __name__ == "__main__":
+    main()