refresh the whole file, making no complicated conversion

Author: WeiguoMa

examples/vqe_qudit_example.py

@@ -1,248 +1,170 @@
-r"""
-You must set the backend explicitly via --backend {jax, tensorflow}.
-AD-based optimization (gradient descent) is enabled for these backends.
-A fallback random-search optimizer is also provided.
-
-Example runs:
-  python vqe_qudit_example.py --backend jax --optimizer gd --dim 3 --layers 2 --steps 200 --lr 0.1 --jit
-  python vqe_qudit_example.py --backend tensorflow --optimizer gd --dim 3 --layers 2 --steps 200 --lr 0.1
-  python vqe_qudit_example.py --backend jax --optimizer random --dim 3 --layers 2 --iters 300
-
-What this script does:
-  - Builds a 2-qudit (d>=3) ansatz with native RY/RZ single-qudit rotations on adjacent levels
-    and an RXX entangler on (0,1) level pairs.
-  - Minimizes the expectation of a simple 2-site Hermitian Hamiltonian:
-        H = N(0) + N(1) + J * [ X_sym(0)\otimes X_sym(1) + Z_sym(0)\otimes Z_sym(1) ]
-    where N = diag(0,1,...,d-1), X_sym = (X + X^\dagger)/2, Z_sym = (Z + Z^\dagger)/2.
 """
+VQE on QuditCircuits.
 
-import argparse
-import math
-import sys
-from typing import Sequence, Tuple
+This example shows how to run a simple VQE on a qudit system using
+`tensorcircuit.QuditCircuit`. We build a compact ansatz using single-qudit
+rotations in selected two-level subspaces and RXX-type entanglers, then
+optimize the energy of a Hermitian "clock–shift" Hamiltonian:
 
-import numpy as np
-import tensorcircuit as tc
-from tensorcircuit.quditcircuit import QuditCircuit
+    H(d) = - J * (X_c \otimes X_c)  -  h * (Z_c \otimes I + I \otimes Z_c)
 
+where, for local dimension `d`,
+- Z_c = (Z + Z^\dagger)/2 is the Hermitian "clock" observable with Z = diag(1, \omega, \omega^2, ..., \omega^{d-1})
+- X_c = (S + S^\dagger)/2 is the Hermitian "shift" observable with S the cyclic shift
+- \omega = exp(2\pi i/d)
 
-# ---------- Hamiltonian helpers ----------
-def symmetrize_hermitian(U: np.ndarray) -> np.ndarray:
-    return 0.5 * (U + U.conj().T)
+The code defaults to a 2-qutrit (d=3) problem but can be changed via CLI flags.
+"""
 
+# import os, sys
+#
+# base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+# if base_dir not in sys.path:
+#     sys.path.insert(0, base_dir)
 
-def build_2site_hamiltonian(
-    d: int, J: float
-) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, float]:
-    N = np.diag(np.arange(d))
-    Xsym = symmetrize_hermitian(tc.backend.numpy(tc.quditgates._x_matrix_func(d)))
-    Zsym = symmetrize_hermitian(tc.backend.numpy(tc.quditgates._z_matrix_func(d)))
-    H0 = N.copy()
-    H1 = N.copy()
-    return H0, H1, Xsym, Zsym, J
+import time
+import argparse
+import tensorcircuit as tc
 
+tc.set_backend("jax")
+tc.set_dtype("complex128")
 
-# ---------- Ansatz ----------
 
+def vqe_forward(param, *, nqudits: int, d: int, nlayers: int, J: float, h: float):
+    """Build a QuditCircuit ansatz and compute ⟨H⟩.
 
-def apply_single_qudit_layer(c: QuditCircuit, qudit: int, thetas: Sequence) -> None:
-    """
-    Apply RY(j,j+1) then RZ(j) for each adjacent level pair.
-    Number of params per site = 2*(d-1).
+    Ansatz:
+      [ for L in 1...nlayers ]
+        - On each site q:
+            RX(q; θ_Lq^(01)) ∘ RY(q; θ_Lq^(12)) ∘ RZ(q; φ_Lq^(0))
+          (subspace indices shown as superscripts)
+        - Entangle neighboring pairs with RXX on subspaces (0,1)
     """
-    d = c._d
-    idx = 0
-    for j, k in [(p, p + 1) for p in range(d - 1)]:
-        c.ry(qudit, theta=thetas[idx], j=j, k=k)
-        idx += 1
-        c.rz(qudit, theta=thetas[idx], j=j)
-        idx += 1
-
-
-def apply_entangler(c: QuditCircuit, theta) -> None:
-    # generalized RXX on (0,1) level pair for both qudits
-    c.rxx(0, 1, theta=theta, j1=0, k1=1, j2=0, k2=1)
-
-
-def build_ansatz(nlayers: int, d: int, params: Sequence) -> QuditCircuit:
-    c = QuditCircuit(2, dim=d)
-    per_site = 2 * (d - 1)
-    per_layer = 2 * per_site + 1  # two sites + entangler
-    assert (
-        len(params) == nlayers * per_layer
-    ), f"params length {len(params)} != {nlayers * per_layer}"
-    off = 0
-    for _ in range(nlayers):
-        th0 = params[off : off + per_site]
-        off += per_site
-        th1 = params[off : off + per_site]
-        off += per_site
-        thE = params[off]
-        off += 1
-        apply_single_qudit_layer(c, 0, th0)
-        apply_single_qudit_layer(c, 1, th1)
-        apply_entangler(c, thE)
-    return c
-
-
-# ---------- Energy ----------
-
-
-def energy_expectation_backend(
-    params_b,
-    d: int,
-    nlayers: int,
-    ham: Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, float],
-):
-    """
-    params_b: 1D backend tensor (jax/tf) of shape [nparams].
-    Returns backend scalar.
-    """
-    bk = tc.backend
-    # Keep differentiability by passing backend scalars into gates
-    plist = [params_b[i] for i in range(params_b.shape[0])]
-    c = build_ansatz(nlayers, d, plist)
-    E = 0.0 + 0.0j
-    H0, H1, Xsym, Zsym, J = ham
-    # Local number operators
-    E = E + c.expectation((tc.gates.Gate(H0), [0]))
-    E = E + c.expectation((tc.gates.Gate(H1), [1]))
-    # Coupling terms as products on separate sites (avoids 9x9 reshaping issues)
-    E = E + J * c.expectation((tc.gates.Gate(Xsym), [0]), (tc.gates.Gate(Xsym), [1]))
-    E = E + J * c.expectation((tc.gates.Gate(Zsym), [0]), (tc.gates.Gate(Zsym), [1]))
-    return bk.real(E)
-
-
-def energy_expectation_numpy(
-    params_np: np.ndarray,
-    d: int,
-    nlayers: int,
-    ham: Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, float],
-) -> float:
-    c = build_ansatz(nlayers, d, params_np.tolist())
-    E = 0.0 + 0.0j
-    H0, H1, Xsym, Zsym, J = ham
-    E += c.expectation((tc.gates.Gate(H0), [0]))
-    E += c.expectation((tc.gates.Gate(H1), [1]))
-    # Coupling terms as products on separate sites (avoids 9x9 reshaping issues)
-    E += J * c.expectation((tc.gates.Gate(Xsym), [0]), (tc.gates.Gate(Xsym), [1]))
-    E += J * c.expectation((tc.gates.Gate(Zsym), [0]), (tc.gates.Gate(Zsym), [1]))
-    return float(np.real(E))
-
-
-# ---------- Optimizers ----------
-
-
-def random_search(fun_numpy, x0_shape, iters=300, seed=42):
-    rng = np.random.default_rng(seed)
-    best_x, best_y = None, float("inf")
-    for _ in range(iters):
-        x = rng.uniform(-math.pi, math.pi, size=x0_shape)
-        y = fun_numpy(x)
-        if y < best_y:
-            best_x, best_y = x, y
-    return best_x, float(best_y)
-
-
-def gradient_descent_ad(energy_bk, x0_np: np.ndarray, steps=200, lr=0.1, jit=False):
-    """
-    energy_bk: (backend_tensor[nparams]) -> backend_scalar
-    Simple gradient descent in numpy space with backend-gradients.
-    """
-    bk = tc.backend
-    if jit:
-        energy_bk = bk.jit(energy_bk)
-    grad_f = bk.grad(energy_bk)
+    if d < 3:
+        raise ValueError("This example assblumes d >= 3 (qutrit or higher).")
+
+    S = tc.quditgates._x_matrix_func(d)
+    Z = tc.quditgates._z_matrix_func(d)
+    Sdag = tc.backend.adjoint(S)
+    Zdag = tc.backend.adjoint(Z)
+
+    c = tc.QuditCircuit(nqudits, dim=d)
 
-    x_np = x0_np.copy()
-    best_x, best_y = x_np.copy(), float("inf")
+    pairs = [(i, i + 1) for i in range(nqudits - 1)]
 
-    def to_np(x):
-        return x if isinstance(x, np.ndarray) else bk.numpy(x)
+    it = iter(param)
 
-    for _ in range(steps):
-        x_b = bk.convert_to_tensor(x_np)  # numpy -> backend tensor
-        g_b = grad_f(x_b)  # backend gradient
-        g = to_np(g_b)  # backend -> numpy
-        x_np = x_np - lr * g  # SGD step in numpy
-        y = float(to_np(energy_bk(bk.convert_to_tensor(x_np))))
-        if y < best_y:
-            best_x, best_y = x_np.copy(), y
-    return best_x, float(best_y)
+    for _ in range(nlayers):
+        for q in range(nqudits):
+            c.rx(q, theta=next(it), j=0, k=1)
+            c.ry(q, theta=next(it), j=1, k=2)
+            c.rz(q, theta=next(it), j=0)
+
+        for i, j in pairs:
+            c.rxx(i, j, theta=next(it), j1=0, k1=1, j2=0, k2=1)
+
+    # H = -J * 1/2 (S_i S_j^\dagger + S_i^\dagger S_j) - h * 1/2 (Z + Z^\dagger)
+    energy = 0.0
+    for i, j in pairs:
+        e_ij = 0.5 * (
+            c.expectation((S, [i]), (Sdag, [j])) + c.expectation((Sdag, [i]), (S, [j]))
+        )
+        energy += -J * tc.backend.real(e_ij)
+    for q in range(nqudits):
+        zq = 0.5 * (c.expectation((Z, [q])) + c.expectation((Zdag, [q])))
+        energy += -h * tc.backend.real(zq)
+    return tc.backend.real(energy)
 
 
-# ---------- CLI ----------
+def build_param_shape(nqudits: int, d: int, nlayers: int):
+    # Per layer per qudit: RX^(01), RY^(12) (or dummy), RZ^(0) = 3 params
+    # Per layer entanglers: len(pairs) parameters
+    pairs = nqudits - 1
+    per_layer = 3 * nqudits + pairs
+    return (nlayers * per_layer,)
 
 
 def main():
-    ap = argparse.ArgumentParser(description="Qudit VQE (explicit backend)")
-    ap.add_argument(
-        "--backend",
-        required=True,
-        choices=["jax", "tensorflow"],
-        help="tensorcircuit backend",
+    parser = argparse.ArgumentParser(
+        description="VQE on QuditCircuit (clock–shift model)"
     )
-    ap.add_argument("--dim", type=int, default=3, help="local qudit dimension d (>=3)")
-    ap.add_argument("--layers", type=int, default=2, help="# ansatz layers")
-    ap.add_argument("--J", type=float, default=0.5, help="coupling strength")
-    ap.add_argument(
-        "--optimizer",
-        type=str,
-        default="gd",
-        choices=["gd", "random"],
-        help="gradient descent (AD) or random search",
+    parser.add_argument(
+        "--d", type=int, default=3, help="Local dimension per site (>=3)"
     )
-    ap.add_argument("--steps", type=int, default=200, help="GD steps")
-    ap.add_argument("--lr", type=float, default=0.1, help="GD learning rate")
-    ap.add_argument("--iters", type=int, default=300, help="random search steps")
-    ap.add_argument("--seed", type=int, default=42, help="RNG seed")
-    ap.add_argument(
-        "--jit",
-        action="store_true",
-        help="enable backend JIT (all backends implement .jit; numpy backend no-ops)",
+    parser.add_argument("--nqudits", type=int, default=2, help="Number of sites")
+    parser.add_argument("--nlayers", type=int, default=3, help="Ansatz depth (layers)")
+    parser.add_argument(
+        "--J", type=float, default=1.0, help="Coupling strength for XcXc term"
     )
-    args = ap.parse_args()
-
-    tc.set_backend(args.backend)
-
-    if args.dim < 3:
-        print("Please use dim >= 3 for qudits.", file=sys.stderr)
-        sys.exit(1)
-
-    d, L = args.dim, args.layers
-    per_layer = 4 * (d - 1) + 1
-    nparams = L * per_layer
-
-    ham = build_2site_hamiltonian(d, args.J)
-
-    print(
-        f"[info] backend={args.backend}, d={d}, layers={L}, params={nparams}, J={args.J}"
+    parser.add_argument(
+        "--h", type=float, default=0.6, help="Field strength for Zc terms"
     )
-
-    if args.optimizer == "random":
-
-        def obj_np(theta_np):
-            return energy_expectation_numpy(theta_np, d, L, ham)
-
-        x, y = random_search(
-            obj_np, x0_shape=(nparams,), iters=args.iters, seed=args.seed
+    parser.add_argument("--steps", type=int, default=200, help="Optimization steps")
+    parser.add_argument("--lr", type=float, default=0.05, help="Learning rate")
+    args = parser.parse_args()
+
+    assert args.d >= 3, "d must be >= 3"
+
+    shape = build_param_shape(args.nqudits, args.d, args.nlayers)
+    param = tc.backend.random_uniform(shape, boundaries=(-0.1, 0.1), seed=42)
+
+    try:
+        import optax
+
+        optimizer = optax.adam(args.lr)
+        vgf = tc.backend.jit(
+            tc.backend.value_and_grad(
+                lambda p: vqe_forward(
+                    p,
+                    nqudits=args.nqudits,
+                    d=args.d,
+                    nlayers=args.nlayers,
+                    J=args.J,
+                    h=args.h,
+                )
+            )
         )
-    else:
-
-        def obj_bk(theta_b):
-            return energy_expectation_backend(theta_b, d, L, ham)
-
-        rng = np.random.default_rng(args.seed)
-        x0 = rng.uniform(-math.pi, math.pi, size=(nparams,))
-        x, y = gradient_descent_ad(
-            obj_bk, x0_np=x0, steps=args.steps, lr=args.lr, jit=args.jit
+        opt_state = optimizer.init(param)
+
+        @tc.backend.jit
+        def train_step(p, opt_state):
+            loss, grads = vgf(p)
+            updates, opt_state = optimizer.update(grads, opt_state, p)
+            p = optax.apply_updates(p, updates)
+            return p, opt_state, loss
+
+        print("Starting VQE optimization (optax/adam)...")
+        loss = None
+        for i in range(args.steps):
+            t0 = time.time()
+            param, opt_state, loss = train_step(param, opt_state)
+            # ensure sync for accurate timing
+            _ = float(loss)
+            if i % 20 == 0:
+                dt = time.time() - t0
+                print(f"Step {i:4d}  loss={loss:.6f}  dt/step={dt:.4f}s")
+        print("Final loss:", float(loss) if loss is not None else "n/a")
+
+    except ModuleNotFoundError:
+        print("Optax not available; using naive gradient descent.")
+        value_and_grad = tc.backend.value_and_grad(
+            lambda p: vqe_forward(
+                p,
+                nqudits=args.nqudits,
+                d=args.d,
+                nlayers=args.nlayers,
+                J=args.J,
+                h=args.h,
+            )
         )
-
-    print("\n=== Result ===")
-    print(f"Energy      : {y:.6f}")
-    print(f"Params shape: {x.shape}")
-    np.set_printoptions(precision=4, suppress=True)
-    print(x[: min(10, x.size)])
+        lr = args.lr
+        loss = None
+        for i in range(args.steps):
+            loss, grads = value_and_grad(param)
+            param = param - lr * grads
+            if i % 20 == 0:
+                print(f"Step {i:4d}  loss={float(loss):.6f}")
+        print("Final loss:", float(loss) if loss is not None else "n/a")
 
 
 if __name__ == "__main__":