Update timeevol.py

change two functions

Author: Huang-Xu-Yang

tensorcircuit/timeevol.py

@@ -427,37 +427,40 @@ def _evol(t: Tensor) -> Tensor:
 ed_evol = hamiltonian_evol
 
 
-@partial(arg_alias, alias_dict={"h_fun": ["hamiltonian"], "t": ["times"]})
-def evol_local(
-    c: Circuit,
-    index: Sequence[int],
-    h_fun: Callable[..., Tensor],
-    t: float,
-    *args: Any,
-    **solver_kws: Any,
-) -> Circuit:
-    """
-    ode evolution of time dependent Hamiltonian on circuit of given indices
-    [only jax backend support for now]
+def _solve_ode(f, s, times, args, solver_kws):
+    rtol = solver_kws.get("rtol", 1e-12)
+    atol = solver_kws.get("atol", 1e-12)
 
-    :param c: _description_
-    :type c: Circuit
-    :param index: qubit sites to evolve
-    :type index: Sequence[int]
-    :param h_fun: h_fun should return a dense Hamiltonian matrix
-        with input arguments time and *args
-    :type h_fun: Callable[..., Tensor]
-    :param t: evolution time
-    :type t: float
-    :return: _description_
-    :rtype: Circuit
-    """
-    s = c.state()
-    n = int(np.log2(s.shape[-1]) + 1e-7)
-    if isinstance(t, float):
-        t = backend.stack([0.0, t])
-    s1 = ode_evol_local(h_fun, s, t, index, None, *args, **solver_kws)
-    return type(c)(n, inputs=s1[-1])
+    ts = backend.convert_to_tensor(times)
+    ts = backend.cast(ts, dtype=rdtypestr)
+
+    if (solver := solver_kws.get("solver", "Dopri5")) == "Dopri5":
+        from jax.experimental.ode import odeint
+        s1 = odeint(f, s, ts, rtol=rtol, atol=atol, *args)
+    
+    else:
+        import diffrax, warnings
+        # Ignore complex warning
+        warnings.simplefilter("ignore", category=UserWarning, append=True)
+        dt0 = solver_kws.get("dt0", 0.01)
+        all_solvers = {"Tsit5": diffrax.Tsit5, "Dopri8": diffrax.Dopri8, "Kvaerno5": diffrax.Kvaerno5}
+        
+        # ODE
+        term = diffrax.ODETerm(lambda t, y, args: f(y, t, *args))
+
+        # solve ODE
+        s1 = diffrax.diffeqsolve(
+        terms = term, 
+        solver = all_solvers[solver](), 
+        t0 = times[0],
+        t1 = times[-1],
+        dt0 = dt0,
+        y0 = s,
+        saveat = diffrax.SaveAt(ts=times),
+        args = args,
+        stepsize_controller =  diffrax.PIDController(rtol=rtol, atol=atol),
+        ).ys
+    return s1
 
 
 def ode_evol_local(
@@ -466,14 +469,17 @@ def ode_evol_local(
     times: Tensor,
     index: Sequence[int],
     callback: Optional[Callable[..., Tensor]] = None,
-    *args: Any,
-    **solver_kws: Any,
+    args: tuple | list = tuple(),
+    solver_kws: dict = dict(),
 ) -> Tensor:
     """
     ODE-based time evolution for a time-dependent Hamiltonian acting on a subsystem of qubits.
 
     This function solves the time-dependent Schrodinger equation using numerical ODE integration.
-    The Hamiltonian is applied only to a specific subset of qubits (indices) in the system.
+    The Hamiltonian is applied only to a specific subset of qubits (indices) in the system. 
+
+    If the solver is 'Dopri5' (default), calls `jax.experimental.ode.odeint`; 
+    otherwise calls `diffrax`.
 
     Note: This function currently only supports the JAX backend.
 
@@ -489,14 +495,17 @@ def ode_evol_local(
     :param callback: Optional function to apply to the state at each time step.
     :type callback: Optional[Callable[..., Tensor]]
     :param args: Additional arguments to pass to the Hamiltonian function.
-    :param solver_kws: Additional keyword arguments to pass to the ODE solver.
+    :type args: tuple | list
+    :param solver_kws: Additional keyword arguments to pass to the ODE solver. 
+        The solver type can be specified: {'Dopri5' (default), 'Tsit5', 'Dopri8', 'Kvaerno5'}.
+        rtol (default: 1e-12) and atol (default: 1e-12) are used to determine how accurately you would like the numerical approximation to your equation.
+        dt0 (default: 0.01) specifies the initial step size.
+    :type solver_kws: dict
     :return: Evolved quantum states at the specified time points. If callback is provided,
         returns the callback results; otherwise returns the state vectors.
     :rtype: Tensor
     """
-    from jax.experimental.ode import odeint
 
-    s = initial_state
     n = int(np.log2(backend.shape_tuple(initial_state)[-1]) + 1e-7)
     l = len(index)
 
@@ -517,47 +526,20 @@ def f(y: Tensor, t: Tensor, *args: Any) -> Tensor:
         y = contractor([y, h], output_edge_order=edges)
         return backend.reshape(y.tensor, [-1])
 
-    ts = backend.convert_to_tensor(times)
-    ts = backend.cast(ts, dtype=rdtypestr)
-    s1 = odeint(f, s, ts, *args, **solver_kws)
-    if not callback:
+    s1 = _solve_ode(f, initial_state, times, args, solver_kws)
+        
+    if  callback is None:
         return s1
-    return backend.stack([callback(s1[i]) for i in range(len(s1))])
-
-
-@partial(arg_alias, alias_dict={"h_fun": ["hamiltonian"], "t": ["times"]})
-def evol_global(
-    c: Circuit, h_fun: Callable[..., Tensor], t: float, *args: Any, **solver_kws: Any
-) -> Circuit:
-    """
-    ode evolution of time dependent Hamiltonian on circuit of all qubits
-    [only jax backend support for now]
-
-    :param c: _description_
-    :type c: Circuit
-    :param h_fun: h_fun should return a **SPARSE** Hamiltonian matrix
-        with input arguments time and *args
-    :type h_fun: Callable[..., Tensor]
-    :param t: _description_
-    :type t: float
-    :return: _description_
-    :rtype: Circuit
-    """
-    s = c.state()
-    n = c._nqubits
-    if isinstance(t, float):
-        t = backend.stack([0.0, t])
-    s1 = ode_evol_global(h_fun, s, t, None, *args, **solver_kws)
-    return type(c)(n, inputs=s1[-1])
+    return backend.stack([callback(a_state) for a_state in s1])
 
 
 def ode_evol_global(
     hamiltonian: Callable[..., Tensor],
     initial_state: Tensor,
     times: Tensor,
     callback: Optional[Callable[..., Tensor]] = None,
-    *args: Any,
-    **solver_kws: Any,
+    args: tuple | list = tuple(),
+    solver_kws: dict = dict(),
 ) -> Tensor:
     """
     ODE-based time evolution for a time-dependent Hamiltonian acting on the entire system.
@@ -566,6 +548,9 @@ def ode_evol_global(
     The Hamiltonian is applied to the full system and should be provided in sparse matrix format
     for efficiency.
 
+    If the solver is 'Dopri5' (default), calls `jax.experimental.ode.odeint`; 
+    otherwise calls `diffrax`.
+
     Note: This function currently only supports the JAX backend.
 
     :param hamiltonian: A function that returns a sparse Hamiltonian matrix for the full system.
@@ -578,25 +563,200 @@ def ode_evol_global(
     :param callback: Optional function to apply to the state at each time step.
     :type callback: Optional[Callable[..., Tensor]]
     :param args: Additional arguments to pass to the Hamiltonian function.
-    :param solver_kws: Additional keyword arguments to pass to the ODE solver.
+    :type args: tuple | list
+    :param solver_kws: Additional keyword arguments to pass to the ODE solver. 
+        The solver type can be specified: {'Dopri5' (default), 'Tsit5', 'Dopri8', 'Kvaerno5'}.
+        rtol (default: 1e-12) and atol (default: 1e-12) are used to determine how accurately you would like the numerical approximation to your equation.
+        dt0 (default: 0.01) specifies the initial step size.
+    :type solver_kws: dict
     :return: Evolved quantum states at the specified time points. If callback is provided,
         returns the callback results; otherwise returns the state vectors.
     :rtype: Tensor
     """
-    from jax.experimental.ode import odeint
-
-    s = initial_state
-    ts = backend.convert_to_tensor(times)
-    ts = backend.cast(ts, dtype=rdtypestr)
 
     def f(y: Tensor, t: Tensor, *args: Any) -> Tensor:
         h = -1.0j * hamiltonian(t, *args)
         return backend.sparse_dense_matmul(h, y)
+    
+    s1 = _solve_ode(f, initial_state, times, args, solver_kws)
 
-    s1 = odeint(f, s, ts, *args, **solver_kws)
-    if not callback:
+    if callback is None:
         return s1
-    return backend.stack([callback(s1[i]) for i in range(len(s1))])
+    return backend.stack([callback(a_state) for a_state in s1])
+
+
+@partial(arg_alias, alias_dict={"h_fun": ["hamiltonian"], "t": ["times"]})
+def evol_local(
+    c: Circuit,
+    index: Sequence[int],
+    h_fun: Callable[..., Tensor],
+    t: float,
+    *args: Any,
+    **solver_kws: Any,
+) -> Circuit:
+    """
+    ode evolution of time dependent Hamiltonian on circuit of given indices
+    [only jax backend support for now]
+
+    :param c: _description_
+    :type c: Circuit
+    :param index: qubit sites to evolve
+    :type index: Sequence[int]
+    :param h_fun: h_fun should return a dense Hamiltonian matrix
+        with input arguments time and *args
+    :type h_fun: Callable[..., Tensor]
+    :param t: evolution time
+    :type t: float
+    :return: _description_
+    :rtype: Circuit
+    """
+    s = c.state()
+    n = int(np.log2(s.shape[-1]) + 1e-7)
+    if isinstance(t, float):
+        t = backend.stack([0.0, t])
+    s1 = ode_evol_local(h_fun, s, t, index, None, *args, **solver_kws)
+    return type(c)(n, inputs=s1[-1])
+
+
+# def ode_evol_local(
+#     hamiltonian: Callable[..., Tensor],
+#     initial_state: Tensor,
+#     times: Tensor,
+#     index: Sequence[int],
+#     callback: Optional[Callable[..., Tensor]] = None,
+#     *args: Any,
+#     **solver_kws: Any,
+# ) -> Tensor:
+#     """
+#     ODE-based time evolution for a time-dependent Hamiltonian acting on a subsystem of qubits.
+
+#     This function solves the time-dependent Schrodinger equation using numerical ODE integration.
+#     The Hamiltonian is applied only to a specific subset of qubits (indices) in the system.
+
+#     Note: This function currently only supports the JAX backend.
+
+#     :param hamiltonian: A function that returns a dense Hamiltonian matrix for the specified
+#         subsystem size. The function signature should be hamiltonian(time, *args) -> Tensor.
+#     :type hamiltonian: Callable[..., Tensor]
+#     :param initial_state: The initial quantum state vector of the full system.
+#     :type initial_state: Tensor
+#     :param times: Time points for which to compute the evolution. Should be a 1D array of times.
+#     :type times: Tensor
+#     :param index: Indices of qubits where the Hamiltonian is applied.
+#     :type index: Sequence[int]
+#     :param callback: Optional function to apply to the state at each time step.
+#     :type callback: Optional[Callable[..., Tensor]]
+#     :param args: Additional arguments to pass to the Hamiltonian function.
+#     :param solver_kws: Additional keyword arguments to pass to the ODE solver.
+#     :return: Evolved quantum states at the specified time points. If callback is provided,
+#         returns the callback results; otherwise returns the state vectors.
+#     :rtype: Tensor
+#     """
+#     from jax.experimental.ode import odeint
+
+#     s = initial_state
+#     n = int(np.log2(backend.shape_tuple(initial_state)[-1]) + 1e-7)
+#     l = len(index)
+
+#     def f(y: Tensor, t: Tensor, *args: Any) -> Tensor:
+#         y = backend.reshape2(y)
+#         y = Gate(y)
+#         h = -1.0j * hamiltonian(t, *args)
+#         h = backend.reshape2(h)
+#         h = Gate(h)
+#         edges = []
+#         for i in range(n):
+#             if i not in index:
+#                 edges.append(y[i])
+#             else:
+#                 j = index.index(i)
+#                 edges.append(h[j])
+#                 h[j + l] ^ y[i]
+#         y = contractor([y, h], output_edge_order=edges)
+#         return backend.reshape(y.tensor, [-1])
+
+#     ts = backend.convert_to_tensor(times)
+#     ts = backend.cast(ts, dtype=rdtypestr)
+#     s1 = odeint(f, s, ts, *args, **solver_kws)
+#     if not callback:
+#         return s1
+#     return backend.stack([callback(s1[i]) for i in range(len(s1))])
+
+
+@partial(arg_alias, alias_dict={"h_fun": ["hamiltonian"], "t": ["times"]})
+def evol_global(
+    c: Circuit, h_fun: Callable[..., Tensor], t: float, *args: Any, **solver_kws: Any
+) -> Circuit:
+    """
+    ode evolution of time dependent Hamiltonian on circuit of all qubits
+    [only jax backend support for now]
+
+    :param c: _description_
+    :type c: Circuit
+    :param h_fun: h_fun should return a **SPARSE** Hamiltonian matrix
+        with input arguments time and *args
+    :type h_fun: Callable[..., Tensor]
+    :param t: _description_
+    :type t: float
+    :return: _description_
+    :rtype: Circuit
+    """
+    s = c.state()
+    n = c._nqubits
+    if isinstance(t, float):
+        t = backend.stack([0.0, t])
+    s1 = ode_evol_global(h_fun, s, t, None, *args, **solver_kws)
+    return type(c)(n, inputs=s1[-1])
+
+
+# def ode_evol_global(
+#     hamiltonian: Callable[..., Tensor],
+#     initial_state: Tensor,
+#     times: Tensor,
+#     callback: Optional[Callable[..., Tensor]] = None,
+#     *args: Any,
+#     **solver_kws: Any,
+# ) -> Tensor:
+#     """
+#     ODE-based time evolution for a time-dependent Hamiltonian acting on the entire system.
+
+#     This function solves the time-dependent Schrodinger equation using numerical ODE integration.
+#     The Hamiltonian is applied to the full system and should be provided in sparse matrix format
+#     for efficiency.
+
+#     Note: This function currently only supports the JAX backend.
+
+#     :param hamiltonian: A function that returns a sparse Hamiltonian matrix for the full system.
+#         The function signature should be hamiltonian(time, *args) -> Tensor.
+#     :type hamiltonian: Callable[..., Tensor]
+#     :param initial_state: The initial quantum state vector.
+#     :type initial_state: Tensor
+#     :param times: Time points for which to compute the evolution. Should be a 1D array of times.
+#     :type times: Tensor
+#     :param callback: Optional function to apply to the state at each time step.
+#     :type callback: Optional[Callable[..., Tensor]]
+#     :param args: Additional arguments to pass to the Hamiltonian function.
+#     :param solver_kws: Additional keyword arguments to pass to the ODE solver.
+#     :return: Evolved quantum states at the specified time points. If callback is provided,
+#         returns the callback results; otherwise returns the state vectors.
+#     :rtype: Tensor
+#     """
+#     from jax.experimental.ode import odeint
+
+#     s = initial_state
+#     ts = backend.convert_to_tensor(times)
+#     ts = backend.cast(ts, dtype=rdtypestr)
+
+#     def f(y: Tensor, t: Tensor, *args: Any) -> Tensor:
+#         h = -1.0j * hamiltonian(t, *args)
+#         return backend.sparse_dense_matmul(h, y)
+
+#     s1 = odeint(f, s, ts, *args, **solver_kws)
+#     if not callback:
+#         return s1
+#     return backend.stack([callback(s1[i]) for i in range(len(s1))])
+
+
 
 
 def chebyshev_evol(