add chebyshev time evolution

Author: refraction-ray

examples/chebyshev_evol.py

@@ -0,0 +1,237 @@
+"""
+Chebyshev time evolution example with CLI options.
+
+This script demonstrates the Chebyshev time evolution method with various options
+for backend, matrix type, and JIT compilation.
+"""
+
+import argparse
+from typing import Any, Tuple
+import numpy as np
+from scipy.linalg import expm
+
+import tensorcircuit as tc
+
+tc.set_dtype("complex128")
+
+
+def create_heisenberg_hamiltonian(num_sites: int, sparse: bool = True) -> Any:
+    """
+    Create Heisenberg Hamiltonian for a 1D chain.
+
+    Args:
+        num_sites: Number of sites in the chain.
+        sparse: Whether to create a sparse matrix.
+
+    Returns:
+        Hamiltonian matrix.
+    """
+    graph = tc.templates.graphs.Line1D(num_sites)
+    return tc.quantum.heisenberg_hamiltonian(graph, sparse=sparse)
+
+
+def create_initial_state(dim: int) -> Any:
+    """
+    Create initial state as equal superposition.
+
+    Args:
+        dim: Dimension of the Hilbert space.
+
+    Returns:
+        Normalized initial state.
+    """
+    psi0 = tc.backend.ones([dim])
+    return psi0 / tc.backend.norm(psi0)
+
+
+def estimate_spectral_bounds(hamiltonian: Any, n_iter: int = 40) -> Tuple[float, float]:
+    """
+    Estimate spectral bounds of the Hamiltonian.
+
+    Args:
+        hamiltonian: Hamiltonian matrix.
+        n_iter: Number of iterations for Lanczos algorithm.
+
+    Returns:
+        Tuple of (E_max, E_min).
+    """
+    print(f"Estimating spectral bounds of Hamiltonian (iterations: {n_iter})...")
+    # Ensure the initial vector is compatible with JAX backend
+    e_max, e_min = tc.timeevol.estimate_spectral_bounds(hamiltonian, n_iter=n_iter)
+    print(f"Estimated result: E_max = {e_max:.4f}, E_min = {e_min:.4f}")
+    return float(e_max), float(e_min)
+
+
+def compare_with_exact_evolution(
+    hamiltonian: Any,
+    initial_state: Any,
+    chebyshev_state: Any,
+    time: float,
+) -> float:
+    """
+    Compare Chebyshev evolution result with exact evolution.
+
+    Args:
+        hamiltonian: Hamiltonian matrix.
+        initial_state: Initial quantum state.
+        chebyshev_state: State evolved with Chebyshev method.
+        time: Evolution time.
+
+    Returns:
+        Fidelity between the two states.
+    """
+    # Exact evolution using matrix exponential
+    if tc.backend.is_sparse(hamiltonian):
+        h = tc.backend.to_dense(hamiltonian)
+    else:
+        h = hamiltonian
+    psi_exact = expm(-1j * np.asarray(h) * time) @ np.asarray(initial_state)
+
+    fidelity = np.abs(np.vdot(psi_exact, np.asarray(chebyshev_state))) ** 2
+    return fidelity
+
+
+def run_chebyshev_evolution(
+    num_sites: int = 8,
+    time: float = 500.0,
+    backend_name: str = "numpy",
+    sparse: bool = True,
+    use_jit: bool = False,
+) -> None:
+    """
+    Run Chebyshev time evolution with specified parameters.
+
+    Args:
+        num_sites: Number of sites in the system.
+        time: Evolution time.
+        backend_name: Backend to use (numpy, jax, tensorflow, pytorch).
+        sparse: Whether to use sparse matrices.
+        use_jit: Whether to use JIT compilation.
+    """
+    # Set backend
+    tc.set_dtype("complex128")  # Ensure dtype is set after backend if needed
+    tc.set_backend(backend_name)
+    backend = tc.backend
+    print(f"Using {backend_name} backend")
+
+    # Create system
+    dim = 2**num_sites
+    graph = tc.templates.graphs.Line1D(num_sites)
+    h_matrix = tc.quantum.heisenberg_hamiltonian(graph, sparse=sparse)
+    print(f"Created Heisenberg Hamiltonian for {num_sites} sites")
+    print(f"Matrix is {'sparse' if sparse else 'dense'}")
+
+    # Create initial state
+    psi0 = create_initial_state(dim)
+    print("Created initial state (equal superposition)")
+
+    # Estimate spectral bounds
+    e_max, e_min = estimate_spectral_bounds(h_matrix, n_iter=40)
+
+    # Prepare Chebyshev evolution function
+    if use_jit:
+        chebyshev_evol_jit = backend.jit(
+            tc.timeevol.chebyshev_evol, static_argnums=(3, 4, 5)
+        )
+        chebyshev_function = chebyshev_evol_jit
+        print("Using JIT compilation")
+    else:
+        chebyshev_function = tc.timeevol.chebyshev_evol
+        print("Not using JIT compilation")
+
+    # Perform Chebyshev evolution
+    print("\nPerforming Chebyshev evolution...")
+    print("--- Testing single time evolution ---")
+    k_estimate = tc.timeevol.estimate_k(time, (e_max, e_min))
+    m_estimate = tc.timeevol.estimate_M(time, (e_max, e_min), k=k_estimate)
+    print(f"Required M (estimated): {m_estimate}")
+    print(f"Required k (estimated): {k_estimate}")
+
+    psi_cheby = chebyshev_function(
+        h_matrix,
+        psi0,
+        t=time,
+        spectral_bounds=(e_max + 0.1, e_min - 0.1),
+        k=k_estimate,
+        M=m_estimate,
+    )
+
+    norm = tc.backend.norm(psi_cheby)
+    print(f"Norm of evolved state: {norm}")
+
+    # Compare with exact evolution
+    print("\nComparing with exact evolution...")
+    fidelity = compare_with_exact_evolution(h_matrix, psi0, psi_cheby, time)
+    print(f"Fidelity for t={time}: {fidelity:.8f}")
+
+
+def main() -> None:
+    """Main function with CLI argument parsing."""
+    parser = argparse.ArgumentParser(
+        description="Chebyshev time evolution example",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Example usage:
+  python chebyshev_evol.py --num_sites 8 --time 500.0
+  python chebyshev_evol.py --backend jax --jit
+  python chebyshev_evol.py --dense --backend jax --jit
+        """,
+    )
+
+    parser.add_argument(
+        "--num_sites",
+        type=int,
+        default=8,
+        help="Number of sites in the system (default: 8)",
+    )
+
+    parser.add_argument(
+        "--time",
+        type=float,
+        default=500.0,
+        help="Evolution time (default: 500.0)",
+    )
+
+    parser.add_argument(
+        "--backend",
+        type=str,
+        default="numpy",
+        choices=["numpy", "jax", "tensorflow", "pytorch"],
+        help="Backend selection (default: numpy)",
+    )
+
+    parser.add_argument(
+        "--dense",
+        dest="sparse",
+        action="store_false",
+        help="Use dense matrices instead of sparse",
+    )
+
+    parser.add_argument(
+        "--sparse",
+        dest="sparse",
+        action="store_true",
+        help="Use sparse matrices (default)",
+    )
+
+    parser.add_argument(
+        "--jit",
+        action="store_true",
+        help="Enable JIT compilation (only works with JAX backend)",
+    )
+
+    parser.set_defaults(sparse=True)
+
+    args = parser.parse_args()
+
+    run_chebyshev_evolution(
+        num_sites=args.num_sites,
+        time=args.time,
+        backend_name=args.backend,
+        sparse=args.sparse,
+        use_jit=args.jit,
+    )
+
+
+if __name__ == "__main__":
+    main()

tensorcircuit/backends/abstract_backend.py

@@ -808,6 +808,21 @@ def solve(self: Any, A: Tensor, b: Tensor, **kws: Any) -> Tensor:
             "Backend '{}' has not implemented `solve`.".format(self.name)
         )
 
+    def special_jv(self: Any, v: int, z: Tensor, M: int) -> Tensor:
+        """
+        Special function: Bessel function of the first kind.
+
+        :param v: The order of the Bessel function.
+        :type v: int
+        :param z: The argument of the Bessel function.
+        :type z: Tensor
+        :return: The value of the Bessel function [J_0, ...J_{v-1}(z)].
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `special_jv`.".format(self.name)
+        )
+
     def searchsorted(self: Any, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
         """
         Find indices where elements should be inserted to maintain order.
@@ -1400,10 +1415,37 @@ def scan(
                 carry = f(carry, x)
 
         return carry
-        # carry = init
-        # for x in xs:
-        #     carry = f(carry, x)
-        # return carry
+
+    def jaxy_scan(
+        self: Any, f: Callable[[Tensor, Tensor], Tensor], init: Tensor, xs: Tensor
+    ) -> Tensor:
+        """
+        This API follows jax scan style. TF use plain for loop
+
+        :param f: _description_
+        :type f: Callable[[Tensor, Tensor], Tensor]
+        :param init: _description_
+        :type init: Tensor
+        :param xs: _description_
+        :type xs: Tensor
+        :raises ValueError: _description_
+        :return: _description_
+        :rtype: Tensor
+        """
+        if xs is None:
+            raise ValueError("Either xs or length must be provided.")
+        if xs is not None:
+            length = len(xs)
+        carry, outputs_to_stack = init, []
+        for i in range(length):
+            if isinstance(xs, (tuple, list)):
+                x = [ele[i] for ele in xs]
+            else:
+                x = xs[i]
+            new_carry, y = f(carry, x)
+            carry = new_carry
+            outputs_to_stack.append(y)
+        return carry, self.stack(outputs_to_stack)
 
     def stop_gradient(self: Any, a: Tensor) -> Tensor:
         """

tensorcircuit/backends/jax_backend.py

@@ -418,6 +418,11 @@ def is_tensor(self, a: Any) -> bool:
     def solve(self, A: Tensor, b: Tensor, assume_a: str = "gen") -> Tensor:  # type: ignore
         return jsp.linalg.solve(A, b, assume_a=assume_a)
 
+    def special_jv(self, v: int, z: Tensor, M: int) -> Tensor:
+        from .jax_ops import bessel_jv_jax_rescaled
+
+        return bessel_jv_jax_rescaled(v, z, M)
+
     def searchsorted(self, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
         if not self.is_tensor(a):
             a = self.convert_to_tensor(a)
@@ -615,6 +620,11 @@ def f_jax(*args: Any, **kws: Any) -> Any:
         carry, _ = libjax.lax.scan(f_jax, init, xs)
         return carry
 
+    def jaxy_scan(
+        self, f: Callable[[Tensor, Tensor], Tensor], init: Tensor, xs: Tensor
+    ) -> Tensor:
+        return libjax.lax.scan(f, init, xs)
+
     def scatter(self, operand: Tensor, indices: Tensor, updates: Tensor) -> Tensor:
         # updates = jnp.reshape(updates, indices.shape)
         # return operand.at[indices].set(updates)