Ejemplo n.º 1
0
def krylovsolve(
    H: Qobj,
    psi0: Qobj,
    tlist: np.array,
    krylov_dim: int,
    e_ops=None,
    options=None,
    progress_bar: bool = None,
    sparse: bool = False,
):
    """
    Time evolution of state vectors for time independent Hamiltonians.
    Evolve the state vector ("psi0") finding an approximation for the time
    evolution operator of Hamiltonian ("H") by obtaining the projection of
    the time evolution operator on a set of small dimensional Krylov
    subspaces (m << dim(H)).

    The output is either the state vector or the expectation values of
    supplied operators ("e_ops") at arbitrary points at ("tlist").

    **Additional options**

    Additional options to krylovsolve can be set with the following:

    * "store_states": stores states even though expectation values are
      requested via the "e_ops" argument.

    * "store_final_state": store final state even though expectation values are
      requested via the "e_ops" argument.

    Parameters
    ----------
    H : :class:`qutip.Qobj`
        System Hamiltonian.
    psi0 : :class: `qutip.Qobj`
        Initial state vector (ket).
    tlist : None / *list* / *array*
        List of times on which to evolve the initial state. If None, nothing
        happens but the code won't break.
    krylov_dim: int
        Dimension of Krylov approximation subspaces used for the time
        evolution approximation.
    e_ops : None / list of :class:`qutip.Qobj`
        Single operator or list of operators for which to evaluate
        expectation values.
    options : Options
        Instance of ODE solver options, as well as krylov parameters.
            atol: controls (approximately) the error desired for the final
                  solution. (Defaults to 1e-8)
            nsteps: maximum number of krylov's internal number of Lanczos
                    iterations. (Defaults to 10000)
    progress_bar : None / BaseProgressBar
         Optional instance of BaseProgressBar, or a subclass thereof, for
         showing the progress of the simulation.
    sparse : bool (default False)
         Use np.array to represent system Hamiltonians. If True, scipy sparse
         arrays are used instead.

    Returns
    -------
    result: :class:`qutip.solver.Result`
        An instance of the class :class:`qutip.solver.Result`, which contains
        either an *array* `result.expect` of expectation values for the times
        `tlist`, or an *array* `result.states` of state vectors corresponding
        to the times `tlist` [if `e_ops` is an empty list].
    """
    # check the physics
    _check_inputs(H, psi0, krylov_dim)

    # check extra inputs
    e_ops, e_ops_dict = _check_e_ops(e_ops)
    pbar = _check_progress_bar(progress_bar)

    # transform inputs type from Qobj to np.ndarray/csr_matrix
    if sparse:
        _H = H.get_data()  # (fast_) csr_matrix
    else:
        _H = H.full().copy()  # np.ndarray
    _psi = psi0.full().copy()
    _psi = _psi / np.linalg.norm(_psi)

    # create internal variable and output containers
    if options is None:
        options = Options(nsteps=10000)
    krylov_results = Result()
    krylov_results.solver = "krylovsolve"

    # handle particular cases of an empty tlist or single element
    n_tlist_steps = len(tlist)
    if n_tlist_steps < 1:
        return krylov_results

    if n_tlist_steps == 1:  # if tlist has only one element, return it
        krylov_results = particular_tlist_or_happy_breakdown(
            tlist, n_tlist_steps, options, psi0, e_ops, krylov_results,
            pbar)  # this will also raise a warning
        return krylov_results

    tf = tlist[-1]
    t0 = tlist[0]

    # optimization step using Lanczos, then reuse it for the first partition
    dim_m = krylov_dim
    krylov_basis, T_m = lanczos_algorithm(_H,
                                          _psi,
                                          krylov_dim=dim_m,
                                          sparse=sparse)

    # check if a happy breakdown occurred
    if T_m.shape[0] < krylov_dim + 1:
        if T_m.shape[0] == 1:
            # this means that the state does not evolve in time, it lies in a
            # symmetry of H subspace. Thus, theres no work to be done.
            krylov_results = particular_tlist_or_happy_breakdown(
                tlist,
                n_tlist_steps,
                options,
                psi0,
                e_ops,
                krylov_results,
                pbar,
                happy_breakdown=True,
            )
            return krylov_results
        else:
            # no optimization is required, convergence is guaranteed.
            delta_t = tf - t0
            n_timesteps = 1
    else:

        # calculate optimal number of internal timesteps.
        delta_t = _optimize_lanczos_timestep_size(T_m,
                                                  krylov_basis=krylov_basis,
                                                  tlist=tlist,
                                                  options=options)
        n_timesteps = int(ceil((tf - t0) / delta_t))

        if n_timesteps >= options.nsteps:
            raise Exception(
                f"Optimization requires a number {n_timesteps} of lanczos iterations, "
                f"which exceeds the defined allowed number {options.nsteps}. This can "
                "be increased via the 'Options.nsteps' property.")

    partitions = _make_partitions(tlist=tlist, n_timesteps=n_timesteps)

    if progress_bar:
        pbar.start(len(partitions))

    # update parameters regarding e_ops
    krylov_results, expt_callback, options, n_expt_op = _e_ops_outputs(
        krylov_results, e_ops, n_tlist_steps, options)

    # parameters for the lazy iteration evolve tlist
    psi_norm = np.linalg.norm(_psi)
    last_t = t0

    for idx, partition in enumerate(partitions):

        evolved_states = _evolve_krylov_tlist(
            H=_H,
            psi0=_psi,
            krylov_dim=dim_m,
            tlist=partition,
            t0=last_t,
            psi_norm=psi_norm,
            krylov_basis=krylov_basis,
            T_m=T_m,
            sparse=sparse,
        )

        if idx == 0:
            krylov_basis = None
            T_m = None
            t_idx = 0

        _psi = evolved_states[-1]
        psi_norm = np.linalg.norm(_psi)
        last_t = partition[-1]

        # apply qobj to each evolved state, remove repeated tail elements
        qobj_evolved_states = [
            Qobj(state, dims=psi0.dims) for state in evolved_states[1:-1]
        ]

        krylov_results = _expectation_values(
            e_ops,
            n_expt_op,
            expt_callback,
            krylov_results,
            qobj_evolved_states,
            partitions,
            idx,
            t_idx,
            options,
        )

        t_idx += len(partition[1:-1])

        pbar.update(idx)

    pbar.finished()

    if e_ops_dict:
        krylov_results.expect = {
            e: krylov_results.expect[n]
            for n, e in enumerate(e_ops_dict.keys())
        }

    return krylov_results