solvers

Solvers

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

Quantum Master Equation solver.

Parameters:

Name	Type	Description	Default
H	Union[Qarray, Callable[[float], Qarray]]	time dependent Hamiltonian function or time-independent Qarray.	required
rho0	Qarray	initial state, must be a density matrix. For statevector evolution, please use sesolve.	required
tlist	Array	time list	required
saveat_tlist	Optional[Array]	list of times at which to save the state. If -1 or [-1], save only at final time. If None, save at all times in tlist. Default: None.	None
c_ops	Optional[Qarray]	qarray list of collapse operators	None
solver_options	Optional[SolverOptions]	SolverOptions with solver options	None

Returns:

Type	Description
Qarray	list of states

Source code in jaxquantum/core/solvers.py

def mesolve(
    H: Union[Qarray, Callable[[float], Qarray]],
    rho0: Qarray,
    tlist: Array,
    saveat_tlist: Optional[Array] = None,
    c_ops: Optional[Qarray] = None,
    solver_options: Optional[SolverOptions] = None,
) -> Qarray:
    """Quantum Master Equation solver.

    Args:
        H: time dependent Hamiltonian function or time-independent Qarray.
        rho0: initial state, must be a density matrix. For statevector evolution, please use sesolve.
        tlist: time list
        saveat_tlist: list of times at which to save the state.
            If -1 or [-1], save only at final time.
            If None, save at all times in tlist. Default: None.
        c_ops: qarray list of collapse operators
        solver_options: SolverOptions with solver options

    Returns:
        list of states
    """

    saveat_tlist = saveat_tlist if saveat_tlist is not None else tlist

    saveat_tlist = jnp.atleast_1d(saveat_tlist)

    c_ops = c_ops if c_ops is not None else Qarray.from_list([])

    # if isinstance(H, Qarray):

    if len(c_ops) == 0 and rho0.qtype != Qtypes.oper:
        logging.warning(
            "Consider using `jqt.sesolve()` instead, as `c_ops` is an empty list and the initial state is not a density matrix."
        )

    ρ0 = rho0.to_dm().to_dense()

    if robust_isscalar(H):
        H = H * identity_like(ρ0)  # treat scalar H as a multiple of the identity

    dims = ρ0.dims
    ρ0 = ρ0.data

    c_ops = c_ops.data

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

    ys = _mesolve_data(Ht_data, ρ0, tlist, saveat_tlist, c_ops,
                       solver_options=solver_options)

    return jnp2jqt(ys, dims=dims)

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

Generate the propagator for a time dependent 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	Optional[Array]	list of times at which to save the state. If -1 or [-1], save only at final time. If None, save at all times in tlist. Default: None.	None

Returns:

Type	Description
	Qarray or List[Qarray]: The propagator for the Hamiltonian at time t. OR a list of propagators for the Hamiltonian at each time in t.

Source code in jaxquantum/core/solvers.py

def propagator(
    H: Union[Qarray, Callable[[float], Qarray]],
    ts: Union[float, Array],
    saveat_tlist: Optional[Array] = None,
    solver_options=None
):
    """ Generate the propagator for a time dependent 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: list of times at which to save the state.
            If -1 or [-1], save only at final time.
            If None, save at all times in tlist. Default: None.

    Returns:
        Qarray or List[Qarray]:
            The propagator for the Hamiltonian at time t.
            OR a list of propagators for the Hamiltonian at each time in t.

    """


    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)
        propagators_data = results.data.squeeze(-1).mT
        propagators = Qarray.create(propagators_data, dims=H_first.space_dims)

        return propagators

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

Schrödinger Equation solver.

Parameters:

Name	Type	Description	Default
H	Union[Qarray, Callable[[float], Qarray]]	time dependent Hamiltonian function or time-independent Qarray.	required
rho0	Qarray	initial state, must be a density matrix. For statevector evolution, please use sesolve.	required
tlist	Array	time list	required
saveat_tlist	Optional[Array]	list of times at which to save the state. If -1 or [-1], save only at final time. If None, save at all times in tlist. Default: None.	None
solver_options	Optional[SolverOptions]	SolverOptions with solver options	None

Returns:

Type	Description
Qarray	list of states

Source code in jaxquantum/core/solvers.py

def sesolve(
    H: Union[Qarray, Callable[[float], Qarray]],
    rho0: Qarray,
    tlist: Array,
    saveat_tlist: Optional[Array] = None,
    solver_options: Optional[SolverOptions] = None,
) -> Qarray:
    """Schrödinger Equation solver.

    Args:
        H: time dependent Hamiltonian function or time-independent Qarray.
        rho0: initial state, must be a density matrix. For statevector evolution, please use sesolve.
        tlist: time list
        saveat_tlist: list of times at which to save the state.
            If -1 or [-1], save only at final time.
            If None, save at all times in tlist. Default: None.
        solver_options: SolverOptions with solver options

    Returns:
        list of states
    """

    saveat_tlist = saveat_tlist if saveat_tlist is not None else tlist

    saveat_tlist = jnp.atleast_1d(saveat_tlist)

    ψ = rho0

    if ψ.qtype == Qtypes.oper:
        raise ValueError(
            "Please use `jqt.mesolve` for initial state inputs in density matrix form."
        )

    ψ = ψ.to_ket().to_dense()

    if robust_isscalar(H):
        H = H * identity_like(ψ)  # treat scalar H as a multiple of the identity

    dims = ψ.dims
    ψ = ψ.data

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

    ys = _sesolve_data(Ht_data, ψ, tlist, saveat_tlist,
                       solver_options=solver_options)

    return jnp2jqt(ys, dims=dims)

solve(f, ρ0, tlist, saveat_tlist, args, solver_options=None)

Gets teh desired solver from diffrax.

Parameters:

Name	Type	Description	Default
f		function defining the ODE	required
ρ0		initial state	required
tlist		time list	required
saveat_tlist		list of times at which to save the state pass in [-1] to save only at final time	required
args		additional arguments to f	required
solver_options	Optional[SolverOptions]	dictionary with solver options	None

Returns:

Type	Description
	solution

Source code in jaxquantum/core/solvers.py

def solve(f, ρ0, tlist, saveat_tlist, args, solver_options: Optional[
    SolverOptions] = None):
    """Gets teh desired solver from diffrax.

    Args:
        f: function defining the ODE
        ρ0: initial state
        tlist: time list
        saveat_tlist: list of times at which to save the state
            pass in [-1] to save only at final time
        args: additional arguments to f
        solver_options: dictionary with solver options

    Returns:
        solution
    """

    # f and ts
    term = ODETerm(f)

    if saveat_tlist.shape[0] == 1 and saveat_tlist == -1:
        saveat = SaveAt(t1=True)
    else:
        saveat = SaveAt(ts=saveat_tlist)

    # solver
    solver_options = solver_options or SolverOptions.create()

    solver_name = solver_options.solver
    solver = getattr(diffrax, solver_name)()
    stepsize_controller = PIDController(rtol=solver_options.rtol, atol=solver_options.atol)

    # solve!
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore",
                                message="Complex dtype support in Diffrax",
                                category=UserWarning)  # NOTE: suppresses complex dtype warning in diffrax
        sol = diffeqsolve(
            term,
            solver,
            t0=tlist[0],
            t1=tlist[-1],
            dt0=tlist[1] - tlist[0],
            y0=ρ0,
            saveat=saveat,
            stepsize_controller=stepsize_controller,
            args=args,
            max_steps=solver_options.max_steps,
            progress_meter=CustomProgressMeter()
            if solver_options.progress_meter
            else NoProgressMeter(),
        )

    return sol