add pauli propogation

Author: refraction-ray

CHANGELOG.md

@@ -16,6 +16,8 @@
 
 - Add multi controller jax support for distrubuted contraction.
 
+- Add `pauli_propagation` function for efficient approximate simulation of Pauli operators in Heisenberg picture.
+
 ### Fixed
 
 - Fix the breaking logic change in jax from dlpack API, dlcapsule -> tensor.

docs/source/api/pauliprop.rst

@@ -0,0 +1,7 @@
+tensorcircuit.pauliprop
+================================================================================
+.. automodule:: tensorcircuit.pauliprop
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :inherited-members:

docs/source/modules.rst

@@ -21,6 +21,7 @@ tensorcircuit
     ./api/mps_base.rst
     ./api/mpscircuit.rst
     ./api/noisemodel.rst
+    ./api/pauliprop.rst
     ./api/quantum.rst
     ./api/quditcircuit.rst
     ./api/quditgates.rst

examples/pauli_propagation_vqe.py

@@ -0,0 +1,120 @@
+"""
+TFIM VQE with Pauli Propagation
+================================
+
+This example demonstrates how to use the Pauli Propagation module to find the
+ground state energy of a Transverse Field Ising Model (TFIM) Hamiltonian.
+
+Pauli Propagation tracks the evolution of Pauli strings in the Heisenberg picture,
+which is particularly efficient for circuits with limited qubit-interaction locality.
+"""
+
+import time
+import numpy as np
+import jax.numpy as jnp
+import optax
+import tensorcircuit as tc
+
+# Use JAX backend for automatic differentiation and efficient scanning
+tc.set_backend("jax")
+
+
+def example_tfim_vqe():
+    # 1. Configuration
+    N = 12  # Number of qubits
+    layers = 4  # Number of circuit layers
+    k = 3  # Local Pauli tracking limit (Pauli string length <= k)
+    g = 1.0  # Transverse field strength
+
+    print(f"TFIM VQE with Pauli Propagation: N={N}, layers={layers}, k={k}, g={g}")
+
+    # 2. Define Hamiltonian: H = \sum Z_i Z_{i+1} + g \sum X_i
+    # We use tc-internal format for efficiency in the engine:
+    # 0 -> I, 1 -> X, 2 -> Y, 3 -> Z
+    num_terms = (N - 1) + N
+    ham_structures = np.zeros((num_terms, N), dtype=int)
+    ham_weights = np.zeros(num_terms, dtype=np.float32)
+
+    idx = 0
+    # Nearest-neighbor ZZ terms
+    for i in range(N - 1):
+        ham_structures[idx, i] = 3  # Z
+        ham_structures[idx, i + 1] = 3  # Z
+        ham_weights[idx] = 1.0
+        idx += 1
+
+    # Transverse field X terms
+    for i in range(N):
+        ham_structures[idx, i] = 1  # X
+        ham_weights[idx] = g
+        idx += 1
+
+    # 3. Parameter Initialization
+    # Each layer has (N-1) Rzz parameters and N Rx parameters
+    params_per_layer = (N - 1) + N
+    total_params = layers * params_per_layer
+    # Using small random initialization
+    params = np.random.normal(scale=0.1, size=total_params).astype(np.float32)
+    params = jnp.array(params)
+
+    # 4. Define Loss Function
+    pp_engine = tc.pauliprop.PauliPropagationEngine(N, k)
+
+    def loss_fn(p_flat):
+        # Reshape parameters for efficient scanning over layers
+        p_layers = jnp.reshape(p_flat, (layers, params_per_layer))
+
+        # Define how one layer of the circuit is constructed
+        def layer_fn(c, p_l):
+            p_rzz = p_l[: N - 1]
+            p_rx = p_l[N - 1 :]
+
+            # 1. Apply Rx gates
+            for i in range(N):
+                c.rx(i, theta=p_rx[i])
+
+            # 2. Apply Rzz gates
+            # Heisenberg ordering is handled automatically by pp_engine
+            # Even-indexed pairs
+            for i in range(0, N - 1, 2):
+                c.rzz(i, i + 1, theta=p_rzz[i // 2])
+            # Odd-indexed pairs
+            num_even = len(range(0, N - 1, 2))
+            for i in range(1, N - 1, 2):
+                c.rzz(i, i + 1, theta=p_rzz[num_even + (i - 1) // 2])
+
+        # Compute expectation by propagating the Hamiltonian back through layers
+        energy = pp_engine.compute_expectation_scan(
+            ham_structures, ham_weights, layer_fn, p_layers
+        )
+        return jnp.real(energy)
+
+    # 5. Optimization Loop
+    optimizer = optax.adam(learning_rate=0.02)
+    opt_state = optimizer.init(params)
+
+    # JIT compile the value and gradient function for maximum performance
+    loss_val_grad_fn = tc.backend.jit(tc.backend.value_and_grad(loss_fn))
+
+    print("\nStarting Optimization Loop...")
+    t_start = time.time()
+    for i in range(201):
+        # Compute value and gradients
+        val, grads = loss_val_grad_fn(params)
+
+        # Apply updates
+        updates, opt_state = optimizer.update(grads, opt_state)
+        params = optax.apply_updates(params, updates)
+
+        if i % 10 == 0:
+            # Sync JAX if measuring time inside the loop is needed
+            # val.block_until_ready()
+            print(f"Step {i:3d}: Energy = {val:.6f}")
+
+    duration = time.time() - t_start
+    print(f"\nOptimization finished in {duration:.2f}s")
+    print(f"Final VQE Energy: {val:.6f}")
+
+
+if __name__ == "__main__":
+    example_tfim_vqe()

llm_experience.md

@@ -78,3 +78,63 @@ This document records specific technical protocols, lessons learned, and advance
 1.  **JAX Vmap in TensorCircuit Wrapper**:
     *   When using `tc.backend.vmap` (alias `K.vmap`) with the JAX backend, do **not** use JAX-native arguments like `in_axes`. The TC wrapper unifies behavior and exposes `vectorized_argnums` (like TensorFlow) instead.
     *   **Protocol**: Always use `vectorized_argnums=(0, 1, ...)` to specify batched arguments, regardless of the backend (JAX/TF/Torch). Passing `in_axes` will raise a `TypeError` because the wrapper function definition doesn't accept it.
+
+## Backend Quirks
+
+1.  **Missing Backend APIs**:
+    *   The `tc.backend` (JAX) interface may lack some convenient NumPy/TensorFlow methods widely used in other backends.
+    *   **Missing**: `K.zeros_like`, `K.ones_like`, `K.any`, `K.minimum`, `K.gather`, `K.outer`.
+    *   **Workarounds**:
+        *   `zeros_like(x)` -> `K.zeros(x.shape, dtype=x.dtype)`
+        *   `probs > 0` (any) -> `K.sum(probs) > 0`
+        *   `minimum(a,b)` -> `K.where(a < b, a, b)`
+        *   `gather(x, idx)` -> `x[idx]`
+        *   `outer(a, b)` -> `a[:, None] * b[None, :]` (Broadcasting)
+
+2.  **JIT Buffer Management**:
+    *   Updating a fixed-size buffer with a dynamic number of new terms in JAX JIT is challenging due to static shape requirements.
+    *   **Pitfall**: `jax.lax.dynamic_update_slice` requires static slice shapes and can behave unexpectedly if logic implies dynamic sizes.
+    *   **Protocol**: Use **masked updates** (`K.where`). Create a mask for valid insertion indices and map source indices to destination indices using modular arithmetic or standard indexing, masked by the valid region. This maintains static graph shapes.
+
+## Pauli Propagation & Operator Evolution (Heisenberg Picture)
+
+1.  **Heisenberg Picture Reverse Order**:
+    *   Pauli Propagation evolves the *observable* rather than the state. 
+    *   **Protocol**: Circuit operations must be applied in **reversed** chronological order (from the measurement gate back to the initial layer). This corresponds to applying the adjoint gate to the operator: $O \to U^\dagger O U$.
+
+2.  **JAX Tracer Accumulation**:
+    *   When initializing an operator state from a Hamiltonian where weights are JAX Tracers (e.g., in VQE optimization), direct indexing and assignment (`state[idx] = w`) will fail with `TracerArrayConversionError`.
+    *   **Protocol**: Use the `at[].add()` or `at[].set()` functional update syntax for compatibility with JIT and AD.
+        ```python
+        state = state.at[target_idx, flat_idx].add(w)
+        ```
+
+3.  **Real-Valued Expectations for Gradients**:
+    *   JAX gradients of loss functions often require the output to be a real-valued scalar. Even if the physics dictates a real expectation value, numerical complex types (even with zero imaginary part) can trigger `TypeError`.
+    *   **Protocol**: Always explicitly take the real part using `K.real()` or `.real` before returning the expectation value from a loss function or engine method.
+
+
+## Module Integration Protocols
+
+1.  **Exporting and Aliasing**:
+    *   Export new modules in `tensorcircuit/__init__.py`. 
+    *   Provide both the module export (e.g. `from . import pauliprop`) and the primary helper function (e.g. `from .pauliprop import pauli_propagation`). 
+    *   Use concise aliases for frequently used functions (e.g., `PauliProp = pauli_propagation`).
+
+2.  **Backwards Compatibility in Helpers**:
+    *   High-level wrapper functions should support multiple input formats (e.g., both raw arrays and convenient list-of-tuples for observables) to be user-friendly while maintaining internal efficiency.
+
+3.  **Standardized Testing Patterns**:
+    *   **Backend Isolation**: Avoid global `tc.set_backend()` in test files. Use the standard fixtures (`npb`, `tfb`, `jaxb`) from `conftest.py` as function arguments.
+    *   **Test Levels**:
+        *   **Unit**: Test initialization, state mapping, and single-gate kernels.
+        *   **Correctness**: Compare results against `tc.Circuit.expectation` or `expectation_ps` for small $N$.
+        *   **AD/Gradients**: Verify that `jax.grad` (or backend equivalent) works on the module's primary interfaces.
+        *   **Scanning**: Verify that loop-optimization interfaces (e.g. `compute_expectation_scan`) match manual application results.
+
+4.  **Documentation & Linting**:
+    *   Achieve **10/10 pylint score** and pass `mypy` before finalizing a module. 
+    *   Follow Google-style docstrings with reStructuredText markers. This is critical for automated documentation generation.
+
+5.  **User Verification (Walkthroughs)**:
+    *   Always provide a production-ready example in `examples/` (e.g., `pauli_propagation_vqe.py`) that showcases a real-world use case (optimization, dynamics, etc.) and demonstrates performance features like JAX JIT and Scanning.

tensorcircuit/__init__.py

@@ -55,6 +55,8 @@
 from . import cloud
 from . import fgs
 from .fgs import FGSSimulator
+from . import pauliprop
+from .pauliprop import pauli_propagation
 from . import timeevol
 
 FGSCircuit = FGSSimulator