cotengra_visualize_path.py

"""
Visualize the contraction path generated by cotengra for a TensorCircuit circuit.
Requires `cotengra` and `matplotlib`.
"""

import cotengra as ctg
import matplotlib.pyplot as plt
import tensorcircuit as tc


def main():
    # 1. Define a quantum circuit
    n = 12
    c = tc.Circuit(n)

    # Apply some layers of gates
    for i in range(n):
        c.H(i)

    for layer in range(4):
        for i in range(n - 1):
            c.cnot(i, i + 1)
        for i in range(n):
            c.rx(i, theta=0.1 * layer)

    print(f"Created a {n}-qubit circuit with {len(c._nodes)} gates.")

    # 2. Extract a stable contraction topology description for cotengra
    # We use the internal `_nodes` which represents the state vector contraction
    info, _ = tc.cons.get_tn_info(c._nodes)
    inputs, output, size_dict = info

    print("Extracted tensor network topology.")
    print(f"Number of input tensors: {len(inputs)}")

    # 3. Set up cotengra optimizer and search for the optimal contraction tree
    print("Searching for an optimal contraction path using cotengra...")
    opt = ctg.ReusableHyperOptimizer(
        methods=["greedy", "kahypar"], max_repeats=32, max_time=2.0, progbar=True
    )

    tree = opt.search(inputs, output, size_dict)

    print("Contraction path search complete.")
    print(f"Contraction cost (FLOPs): {tree.contraction_cost():.2e}")
    print(f"Maximum memory (sizes): {tree.max_size():.2e}")

    # 4. Visualize the contraction tree
    print("Visualizing the contraction tree...")
    _, axes = plt.subplots(1, 2, figsize=(16, 7))

    # Plot the ring visualization
    tree.plot_ring(ax=axes[0])
    axes[0].set_title("Ring Tree Visualization")

    # Plot the tent visualization
    tree.plot_tent(ax=axes[1])
    axes[1].set_title("Tent Tree Visualization")

    # Save the figure
    output_filename = "cotengra_contraction_path.png"
    plt.savefig(output_filename, bbox_inches="tight", dpi=300)
    print(f"Visualization saved to {output_filename}")


if __name__ == "__main__":
    main()