solvers

CustomProgressMeter

Bases: TqdmProgressMeter

JAXQuantum's default Diffrax progress bar.

Source code in jaxquantum/core/solvers.py

class CustomProgressMeter(diffrax.TqdmProgressMeter):
    """JAXQuantum's default Diffrax progress bar."""

    @staticmethod
    def _init_bar() -> tqdm.tqdm:
        bar_format = (
            "{desc}: {percentage:3.0f}% |{bar}| "
            "[{elapsed}<{remaining}, {rate_fmt}{postfix}]"
        )
        return tqdm.tqdm(
            total=100, bar_format=bar_format, unit="%", colour="MAGENTA", ascii="░▒█"
        )

SolverOptions

Options forwarded to :func:diffrax.diffeqsolve.

Attributes:

Name	Type	Description
solver	AbstractSolver | str	Native Diffrax solver; strings are deprecated.
stepsize_controller	AbstractStepSizeController | str	Native Diffrax controller; strings are deprecated.
stepsize_controller_kwargs	dict[str, Any] | None	Deprecated controller constructor arguments.
saveat	SaveAt | None	Custom save policy; overrides JAXQuantum's save-time handling.
dt0	float | Array | None | Literal['tlist']	Initial step, "tlist" for the first interval, or None for Diffrax's automatic choice.
adjoint	AbstractAdjoint | None	Differentiation strategy. None uses Diffrax's default.
event	Event | None	Native Diffrax termination event.
max_steps	int | None	Maximum solver steps.
throw	bool	Whether unsuccessful solves raise an exception.
progress_meter	bool | Literal['default'] | AbstractProgressMeter | None	None, "default", or a native progress meter. Booleans are deprecated.
solver_state	Any	Solver state used to continue a previous solve.
controller_state	Any	Controller state used to continue a previous solve.
made_jump	bool | Array | None	Previous jump state used when continuing a solve.

saveat=None saves at saveat_tlist (or tlist when omitted). adjoint=None and progress_meter=None preserve Diffrax's defaults. Native Diffrax objects pass through unchanged. Legacy values still work and issue a FutureWarning.

Source code in jaxquantum/core/solvers.py

@struct.dataclass
class SolverOptions:
    """Options forwarded to :func:`diffrax.diffeqsolve`.

    Attributes:
        solver: Native Diffrax solver; strings are deprecated.
        stepsize_controller: Native Diffrax controller; strings are deprecated.
        stepsize_controller_kwargs: Deprecated controller constructor arguments.
        saveat: Custom save policy; overrides JAXQuantum's save-time handling.
        dt0: Initial step, ``"tlist"`` for the first interval, or ``None`` for
            Diffrax's automatic choice.
        adjoint: Differentiation strategy. ``None`` uses Diffrax's default.
        event: Native Diffrax termination event.
        max_steps: Maximum solver steps.
        throw: Whether unsuccessful solves raise an exception.
        progress_meter: ``None``, ``"default"``, or a native progress meter.
            Booleans are deprecated.
        solver_state: Solver state used to continue a previous solve.
        controller_state: Controller state used to continue a previous solve.
        made_jump: Previous jump state used when continuing a solve.

    ``saveat=None`` saves at ``saveat_tlist`` (or ``tlist`` when omitted).
    ``adjoint=None`` and ``progress_meter=None`` preserve Diffrax's defaults.
    Native Diffrax objects pass through unchanged. Legacy values still work and
    issue a ``FutureWarning``.
    """

    progress_meter: bool | Literal["default"] | diffrax.AbstractProgressMeter | None = (
        struct.field(pytree_node=False, default="default")
    )
    solver: diffrax.AbstractSolver | str = struct.field(
        pytree_node=False, default_factory=diffrax.Tsit5
    )
    max_steps: int | None = struct.field(pytree_node=False, default=100_000)
    stepsize_controller: diffrax.AbstractStepSizeController | str = struct.field(
        pytree_node=False, default_factory=_default_stepsize_controller
    )
    stepsize_controller_kwargs: dict[str, Any] | None = struct.field(
        pytree_node=False, default=None
    )
    saveat: diffrax.SaveAt | None = struct.field(pytree_node=False, default=None)
    dt0: float | Array | None | Literal["tlist"] = struct.field(
        pytree_node=False, default="tlist"
    )
    adjoint: diffrax.AbstractAdjoint | None = struct.field(
        pytree_node=False, default=None
    )
    event: diffrax.Event | None = struct.field(pytree_node=False, default=None)
    throw: bool = struct.field(pytree_node=False, default=True)
    solver_state: Any = None
    controller_state: Any = None
    made_jump: bool | Array | None = None

    @classmethod
    def create(
        cls,
        progress_meter: bool = True,
        solver: str = "Tsit5",
        max_steps: int = 100_000,
        stepsize_controller: str = "PIDController",
        stepsize_controller_kwargs: dict[str, Any] | None = None,
    ) -> "SolverOptions":
        """Create options with the deprecated string-based interface."""
        warnings.warn(
            "SolverOptions.create() is deprecated; use SolverOptions with native "
            "objects, such as solver=diffrax.Tsit5() and "
            "progress_meter='default'.",
            FutureWarning,
            stacklevel=2,
        )
        return cls(
            solver=_diffrax_object(solver, diffrax.AbstractSolver),
            stepsize_controller=_diffrax_object(
                stepsize_controller,
                diffrax.AbstractStepSizeController,
                _legacy_controller_kwargs(
                    stepsize_controller, stepsize_controller_kwargs
                ),
            ),
            max_steps=max_steps,
            progress_meter="default" if progress_meter else None,
        )

create(progress_meter=True, solver='Tsit5', max_steps=100000, stepsize_controller='PIDController', stepsize_controller_kwargs=None) classmethod

Create options with the deprecated string-based interface.

Source code in jaxquantum/core/solvers.py

@classmethod
def create(
    cls,
    progress_meter: bool = True,
    solver: str = "Tsit5",
    max_steps: int = 100_000,
    stepsize_controller: str = "PIDController",
    stepsize_controller_kwargs: dict[str, Any] | None = None,
) -> "SolverOptions":
    """Create options with the deprecated string-based interface."""
    warnings.warn(
        "SolverOptions.create() is deprecated; use SolverOptions with native "
        "objects, such as solver=diffrax.Tsit5() and "
        "progress_meter='default'.",
        FutureWarning,
        stacklevel=2,
    )
    return cls(
        solver=_diffrax_object(solver, diffrax.AbstractSolver),
        stepsize_controller=_diffrax_object(
            stepsize_controller,
            diffrax.AbstractStepSizeController,
            _legacy_controller_kwargs(
                stepsize_controller, stepsize_controller_kwargs
            ),
        ),
        max_steps=max_steps,
        progress_meter="default" if progress_meter else None,
    )

mesolve(H, rho0, tlist, saveat_tlist=None, c_ops=None, solver_options=None)

Solve a Lindblad master equation and return the saved states.

Parameters:

Name	Type	Description	Default
H	Qarray | Callable[[float], Qarray]	Static Hamiltonian or callable H(t).	required
rho0	Qarray	Initial ket or density matrix.	required
tlist	Array	Integration interval; also the default save times.	required
saveat_tlist	Array | None	Save times. An empty array saves only the final state.	None
c_ops	Qarray | None	Collapse operators.	None
solver_options	SolverOptions | None	Native Diffrax configuration.	None

Returns:

Type	Description
Qarray	Saved density matrices as a batched Qarray.

See Also

:func:mesolve_result returns the complete Diffrax solution.

Source code in jaxquantum/core/solvers.py

def mesolve(
    H: Qarray | Callable[[float], Qarray],
    rho0: Qarray,
    tlist: Array,
    saveat_tlist: Array | None = None,
    c_ops: Qarray | None = None,
    solver_options: SolverOptions | None = None,
) -> Qarray:
    """Solve a Lindblad master equation and return the saved states.

    Args:
        H: Static Hamiltonian or callable ``H(t)``.
        rho0: Initial ket or density matrix.
        tlist: Integration interval; also the default save times.
        saveat_tlist: Save times. An empty array saves only the final state.
        c_ops: Collapse operators.
        solver_options: Native Diffrax configuration.

    Returns:
        Saved density matrices as a batched ``Qarray``.

    See Also:
        :func:`mesolve_result` returns the complete Diffrax solution.
    """
    solution = mesolve_result(
        H,
        rho0,
        tlist,
        saveat_tlist=saveat_tlist,
        c_ops=c_ops,
        solver_options=solver_options,
    )
    qdims = Qdims((rho0.space_dims, rho0.space_dims))
    return Qarray._from_impl(DenseImpl._make(solution.ys), qdims)

mesolve_result(H, rho0, tlist, saveat_tlist=None, c_ops=None, solver_options=None)

Solve a Lindblad master equation and return its Diffrax solution.

Use this form for solver statistics, events, dense interpolation, custom SaveAt functions, or continuation state.

Source code in jaxquantum/core/solvers.py

def mesolve_result(
    H: Qarray | Callable[[float], Qarray],
    rho0: Qarray,
    tlist: Array,
    saveat_tlist: Array | None = None,
    c_ops: Qarray | None = None,
    solver_options: SolverOptions | None = None,
) -> diffrax.Solution:
    """Solve a Lindblad master equation and return its Diffrax solution.

    Use this form for solver statistics, events, dense interpolation, custom
    ``SaveAt`` functions, or continuation state.
    """
    collapse_ops = c_ops if c_ops is not None else Qarray.from_list([])

    if len(collapse_ops) == 0 and rho0.qtype != Qtypes.oper:
        logging.warning(  # noqa: LOG015
            "Consider sesolve(): no collapse operators were provided and the "
            "initial state is not a density matrix."
        )

    rho_data = rho0.to_dm().to_dense()

    if robust_isscalar(H):
        H = H * identity_like(rho_data)

    if isinstance(H, Qarray):
        H_data = lambda t: H.data
    else:
        H_data = lambda t: H(t).data

    return _mesolve_result_data(
        H_data,
        rho_data.data,
        tlist,
        saveat_tlist,
        collapse_ops.data,
        solver_options=solver_options,
    )

propagator(H, ts, saveat_tlist=None, solver_options=None)

Generate a propagator for a Hamiltonian.

Parameters:

Name	Type	Description	Default
H	Qarray or callable	A Qarray static Hamiltonian OR a function that takes a time argument and returns a Hamiltonian.	required
ts	float or Array	A single time point or an Array of time points.	required
saveat_tlist	Array | None	Times at which to save the propagator.	None
solver_options	SolverOptions | None	Native Diffrax configuration for time-dependent input.	None

Returns:

Type	Description
Qarray	The propagator at each saved time.

Source code in jaxquantum/core/solvers.py

def propagator(
    H: Qarray | Callable[[float], Qarray],
    ts: float | Array,
    saveat_tlist: Array | None = None,
    solver_options: SolverOptions | None = None,
) -> Qarray:
    """Generate a propagator for a Hamiltonian.

    Args:
        H (Qarray or callable):
            A Qarray static Hamiltonian OR
            a function that takes a time argument and returns a Hamiltonian.
        ts (float or Array):
            A single time point or
            an Array of time points.
        saveat_tlist: Times at which to save the propagator.
        solver_options: Native Diffrax configuration for time-dependent input.

    Returns:
        The propagator at each saved time.
    """
    ts_is_scalar = robust_isscalar(ts)
    H_is_qarray = isinstance(H, Qarray)

    if H_is_qarray:
        return (-1j * H * ts).expm()
    else:
        if ts_is_scalar:
            H_first = H(0.0)
            if ts == 0:
                return identity_like(H_first)
            ts = jnp.array([0.0, ts])
        else:
            H_first = H(ts[0])

        basis_states = multi_mode_basis_set(H_first.space_dims)
        results = sesolve(
            H,
            basis_states,
            ts,
            saveat_tlist=saveat_tlist,
            solver_options=solver_options,
        )
        # results.data is (T, M, M): T times, M evolved basis kets (batch), M
        # ket components. Transpose the last two axes so each time slice is a
        # propagator whose columns are the evolved basis states. No squeeze:
        # kets no longer carry a trailing singleton.
        propagators_data = results.data.mT
        return Qarray.create(propagators_data, dims=H_first.space_dims)

sesolve(H, rho0, tlist, saveat_tlist=None, solver_options=None)

Solve a Schrödinger equation and return the saved states.

Parameters:

Name	Type	Description	Default
H	Qarray | Callable[[float], Qarray]	Static Hamiltonian or callable H(t).	required
rho0	Qarray	Initial ket.	required
tlist	Array	Integration interval; also the default save times.	required
saveat_tlist	Array | None	Save times. An empty array saves only the final state.	None
solver_options	SolverOptions | None	Native Diffrax configuration.	None

Returns:

Type	Description
Qarray	Saved kets as a batched Qarray.

See Also

:func:sesolve_result returns the complete Diffrax solution.

Source code in jaxquantum/core/solvers.py

def sesolve(
    H: Qarray | Callable[[float], Qarray],
    rho0: Qarray,
    tlist: Array,
    saveat_tlist: Array | None = None,
    solver_options: SolverOptions | None = None,
) -> Qarray:
    """Solve a Schrödinger equation and return the saved states.

    Args:
        H: Static Hamiltonian or callable ``H(t)``.
        rho0: Initial ket.
        tlist: Integration interval; also the default save times.
        saveat_tlist: Save times. An empty array saves only the final state.
        solver_options: Native Diffrax configuration.

    Returns:
        Saved kets as a batched ``Qarray``.

    See Also:
        :func:`sesolve_result` returns the complete Diffrax solution.
    """
    solution = sesolve_result(
        H,
        rho0,
        tlist,
        saveat_tlist=saveat_tlist,
        solver_options=solver_options,
    )
    return Qarray._from_impl(DenseImpl._make(solution.ys), rho0.qdims)

sesolve_result(H, rho0, tlist, saveat_tlist=None, solver_options=None)

Solve a Schrödinger equation and return its Diffrax solution.

Use this form for solver statistics, events, dense interpolation, custom SaveAt functions, or continuation state.

Source code in jaxquantum/core/solvers.py

def sesolve_result(
    H: Qarray | Callable[[float], Qarray],
    rho0: Qarray,
    tlist: Array,
    saveat_tlist: Array | None = None,
    solver_options: SolverOptions | None = None,
) -> diffrax.Solution:
    """Solve a Schrödinger equation and return its Diffrax solution.

    Use this form for solver statistics, events, dense interpolation, custom
    ``SaveAt`` functions, or continuation state.
    """
    if rho0.qtype == Qtypes.oper:
        raise ValueError("Use mesolve() for an initial density matrix.")

    state = rho0.to_ket().to_dense()

    if robust_isscalar(H):
        H = H * identity_like(state)

    if isinstance(H, Qarray):
        H_data = lambda t: H.data
    else:
        H_data = lambda t: H(t).data

    return _sesolve_result_data(
        H_data,
        state.data,
        tlist,
        saveat_tlist,
        solver_options=solver_options,
    )

solve(f, y0, tlist, saveat_tlist=None, args=None, solver_options=None)

Solve an ODE using native Diffrax configuration from SolverOptions.

Source code in jaxquantum/core/solvers.py

def solve(
    f: Callable,
    y0: Array,
    tlist: Array,
    saveat_tlist: Array | None = None,
    args: Any = None,
    solver_options: SolverOptions | None = None,
) -> diffrax.Solution:
    """Solve an ODE using native Diffrax configuration from ``SolverOptions``."""
    options = SolverOptions() if solver_options is None else solver_options
    if _uses_legacy_options(options):
        warnings.warn(
            "String and boolean SolverOptions values are deprecated; use native "
            "objects such as solver=diffrax.Tsit5(), "
            "stepsize_controller=diffrax.PIDController(...), and "
            "progress_meter='default' or None.",
            FutureWarning,
            stacklevel=2,
        )
    kwargs = {
        "saveat": _resolve_saveat(options, tlist, saveat_tlist),
        "stepsize_controller": _resolve_stepsize_controller(options),
        "args": args,
        "max_steps": options.max_steps,
        "throw": options.throw,
    }
    optional = {
        "adjoint": options.adjoint,
        "event": options.event,
        "progress_meter": _resolve_progress_meter(options.progress_meter),
        "solver_state": options.solver_state,
        "controller_state": options.controller_state,
        "made_jump": options.made_jump,
    }
    kwargs.update(
        (name, value) for name, value in optional.items() if value is not None
    )

    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            message="Complex dtype support in Diffrax",
            category=UserWarning,
        )
        return diffrax.diffeqsolve(
            diffrax.ODETerm(f),
            _resolve_solver(options.solver),
            t0=tlist[0],
            t1=tlist[-1],
            dt0=_resolve_dt0(options, tlist),
            y0=y0,
            **kwargs,
        )