示例#1
0
    def run(self, rho0, tlist, e_ops=None, ado_init=False, ado_return=False):
        """
        Solve for the time evolution of the system.

        Parameters
        ----------
        rho0 : Qobj or HierarchyADOsState or numpy.array
            Initial state (:obj:`~Qobj` density matrix) of the system
            if ``ado_init`` is ``False``.

            If ``ado_init`` is ``True``, then ``rho0`` should be an
            instance of :obj:`~HierarchyADOsState` or a numpy array
            giving the initial state of all ADOs. Usually
            the state of the ADOs would be determine from a previous call
            to ``.run(..., ado_return=True)``. For example,
            ``result = solver.run(..., ado_return=True)`` could be followed
            by ``solver.run(result.ado_states[-1], tlist, ado_init=True)``.

            If a numpy array is passed its shape must be
            ``(number_of_ados, n, n)`` where ``(n, n)`` is the system shape
            (i.e. shape of the system density matrix) and the ADOs must be
            in the same order as in ``.ados.labels``.

        tlist : list
            An ordered list of times at which to return the value of the state.

        e_ops : Qobj / callable / list / dict / None, optional
            A list or dictionary of operators as `Qobj` and/or callable
            functions (they can be mixed) or a single operator or callable
            function. For an operator ``op``, the result will be computed
            using ``(state * op).tr()`` and the state at each time ``t``. For
            callable functions, ``f``, the result is computed using
            ``f(t, ado_state)``. The values are stored in ``expect`` on
            (see the return section below).

        ado_init: bool, default False
            Indicates if initial condition is just the system state, or a
            numpy array including all ADOs.

        ado_return: bool, default True
            Whether to also return as output the full state of all ADOs.

        Returns
        -------
        :class:`qutip.solver.Result`
            The results of the simulation run, with the following attributes:

            * ``times``: the times ``t`` (i.e. the ``tlist``).

            * ``states``: the system state at each time ``t`` (only available
              if ``e_ops`` was ``None`` or if the solver option
              ``store_states`` was set to ``True``).

            * ``ado_states``: the full ADO state at each time (only available
              if ``ado_return`` was set to ``True``). Each element is an
              instance of :class:`HierarchyADOsState`.            .
              The state of a particular ADO may be extracted from
              ``result.ado_states[i]`` by calling :meth:`.extract`.

            * ``expect``: the value of each ``e_ops`` at time ``t`` (only
              available if ``e_ops`` were given). If ``e_ops`` was passed
              as a dictionary, then ``expect`` will be a dictionary with
              the same keys as ``e_ops`` and values giving the list of
              outcomes for the corresponding key.
        """
        e_ops, expected = self._convert_e_ops(e_ops)
        e_ops_callables = any(not isinstance(op, Qobj)
                              for op in e_ops.values())

        n = self._sys_shape
        rho_shape = (n, n)
        rho_dims = self._sys_dims
        hierarchy_shape = (self._n_ados, n, n)

        output = Result()
        output.solver = "HEOMSolver"
        output.times = tlist
        if e_ops:
            output.expect = expected
        if not e_ops or self.options.store_states:
            output.states = []

        if ado_init:
            if isinstance(rho0, HierarchyADOsState):
                rho0_he = rho0._ado_state
            else:
                rho0_he = rho0
            if rho0_he.shape != hierarchy_shape:
                raise ValueError(
                    f"ADOs passed with ado_init have shape {rho0_he.shape}"
                    f"but the solver hierarchy shape is {hierarchy_shape}")
            rho0_he = rho0_he.reshape(n**2 * self._n_ados)
        else:
            rho0_he = np.zeros([n**2 * self._n_ados], dtype=complex)
            rho0_he[:n**2] = rho0.full().ravel('F')

        if ado_return:
            output.ado_states = []

        solver = self._ode
        solver.set_initial_value(rho0_he, tlist[0])

        self.progress_bar.start(len(tlist))
        for t_idx, t in enumerate(tlist):
            self.progress_bar.update(t_idx)
            if t_idx != 0:
                solver.integrate(t)
                if not solver.successful():
                    raise RuntimeError(
                        "HEOMSolver ODE integration error. Try increasing"
                        " the nsteps given in the HEOMSolver options"
                        " (which increases the allowed substeps in each"
                        " step between times given in tlist).")

            rho = Qobj(
                solver.y[:n**2].reshape(rho_shape, order='F'),
                dims=rho_dims,
            )
            if self.options.store_states:
                output.states.append(rho)
            if ado_return or e_ops_callables:
                ado_state = HierarchyADOsState(
                    rho, self.ados, solver.y.reshape(hierarchy_shape))
            if ado_return:
                output.ado_states.append(ado_state)
            for e_key, e_op in e_ops.items():
                if isinstance(e_op, Qobj):
                    e_result = (rho * e_op).tr()
                else:
                    e_result = e_op(t, ado_state)
                output.expect[e_key].append(e_result)

        self.progress_bar.finished()
        return output