Skip to content

dq.sesolve

sesolve(
    H: QArrayLike | TimeQArray,
    psi0: QArrayLike,
    tsave: ArrayLike,
    *,
    exp_ops: list[QArrayLike] | None = None,
    method: Method = Tsit5(),
    gradient: Gradient | None = None,
    save_states: bool = True,
    cartesian_batching: bool = True,
    progress_meter: AbstractProgressMeter | bool | None = None,
    t0: ScalarLike | None = None,
    save_extra: Callable[[Array], PyTree] | None = None
) -> SESolveResult

Solve the Schrödinger equation.

This function computes the evolution of the state vector \(\ket{\psi(t)}\) at time \(t\), starting from an initial state \(\ket{\psi_0}\), according to the Schrödinger equation (with \(\hbar=1\) and where time is implicit(1)) $$ \frac{\dd\ket{\psi}}{\dt} = -i H \ket{\psi}, $$ where \(H\) is the system's Hamiltonian.

  1. With explicit time dependence:
    • \(\ket\psi\to\ket{\psi(t)}\)
    • \(H\to H(t)\)

Parameters:

  • H (qarray-like or timeqarray of shape (...H, n, n)) –

    Hamiltonian.

  • psi0 (qarray-like of shape (...psi0, n, 1)) –

    Initial state.

  • tsave (array-like of shape (ntsave,)) –

    Times at which the states and expectation values are saved. The equation is solved from tsave[0] to tsave[-1], or from t0 to tsave[-1] if t0 is specified.

  • exp_ops (list of qarray-like, each of shape (n, n)) –

    List of operators for which the expectation value is computed.

  • method –

    Method for the integration. Defaults to dq.method.Tsit5 (supported: Tsit5, Dopri5, Dopri8, Kvaerno3, Kvaerno5, Euler, Expm).

  • gradient –

    Algorithm used to compute the gradient. The default is method-dependent, refer to the documentation of the chosen method for more details.

Returns:

  • dq.SESolveResult object holding the result of the Schrödinger equation integration. Use result.states to access the saved states and result.expects to access the saved expectation values.

    Detailed result API
    dq.SESolveResult
    

    Attributes:

    • states (qarray of shape (..., nsave, n, 1)) - Saved states with nsave = ntsave, or nsave = 1 if save_states=False.
    • final_state (qarray of shape (..., n, 1)) - Saved final state.
    • expects (array of shape (..., len(exp_ops), ntsave) or None) - Saved expectation values, if specified by exp_ops.
    • extra (PyTree or None) - Extra data saved with save_extra() if specified.
    • infos (PyTree or None) - Method-dependent information on the resolution.
    • tsave (array of shape (ntsave,)) - Times for which results were saved.
    • method (Method) - Method used.
    • gradient (Gradient) - Gradient used.
    • options (Options) - Options used.

Other Parameters:

  • save_states –

    If True, the state is saved at every time in tsave, otherwise only the final state is returned. Defaults to True.

  • cartesian_batching –

    If True, batched arguments are treated as separated batch dimensions, otherwise the batching is performed over a single shared batch dimension. Defaults to True.

  • progress_meter –

    Progress meter indicating how far the solve has progressed. Defaults to None which uses the global default progress meter (see dq.set_progress_meter()). Set to True for a tqdm progress meter, and False for no output. If gradients are computed, the progress meter only displays during the forward pass.

  • t0 –

    Initial time. If None, defaults to the first time in tsave. Defaults to None.

  • save_extra –

    A function with signature f(QArray) -> PyTree that takes a state as input and returns a PyTree. This can be used to save additional arbitrary data during the integration, accessible in result.extra. Defaults to None.

Examples:

import dynamiqs as dq
import jax.numpy as jnp

n = 16
a = dq.destroy(n)

H = a.dag() @ a
psi0 = dq.coherent(n, 1.0)
tsave = jnp.linspace(0, 1.0, 11)

result = dq.sesolve(H, psi0, tsave)
print(result)
Output
|██████████| 100.0% ◆ elapsed 1.50ms ◆ remaining 0.00ms
==== SESolveResult ====
Method : Tsit5
Infos  : 11 steps (11 accepted, 0 rejected)
States : QArray complex64 (11, 16, 1) | 1.4 Kb

Advanced use-cases

Defining a time-dependent Hamiltonian

If the Hamiltonian depends on time, it can be converted to a timeqarray using dq.pwc(), dq.modulated(), or dq.timecallable(). See the Time-dependent operators tutorial for more details.

Running multiple simulations concurrently

Both the Hamiltonian H and the initial state psi0 can be batched to solve multiple Schrödinger equations concurrently. All other arguments are common to every batch. The resulting states and expectation values are batched according to the leading dimensions of H and psi0. The behaviour depends on the value of the cartesian_batching option.

The results leading dimensions are

... = ...H, ...psi0
For example if:

  • H has shape (2, 3, n, n),
  • psi0 has shape (4, n, 1),

then result.states has shape (2, 3, 4, ntsave, n, 1).

The results leading dimensions are

... = ...H = ...psi0  # (once broadcasted)
For example if:

  • H has shape (2, 3, n, n),
  • psi0 has shape (3, n, 1),

then result.states has shape (2, 3, ntsave, n, 1).

See the Batching simulations tutorial for more details.