<commit_before>"""
This module provides solvers for the unitary Schrodinger equation.
"""

__all__ = ['sesolve']

import os
import types
import numpy as np
import scipy.integrate
from scipy.linalg import norm as la_norm
from qutip.cy.stochastic import normalize_inplace
import qutip.settings as qset
from qutip.qobj import Qobj
from qutip.qobjevo import QobjEvo
from qutip.cy.spconvert import dense1D_to_fastcsr_ket, dense2D_to_fastcsr_fmode
from qutip.cy.spmatfuncs import (cy_expect_psi, cy_ode_psi_func_td,
                                cy_ode_psi_func_td_with_state)
from qutip.solver import Result, Options, config, solver_safe, SolverSystem
from qutip.superoperator import vec2mat
from qutip.ui.progressbar import (BaseProgressBar, TextProgressBar)
from qutip.cy.openmp.utilities import check_use_openmp, openmp_components

def sesolve(H, psi0, tlist, e_ops=None, args=None, options=None,
            progress_bar=None, _safe_mode=True):
    """
    Schrödinger equation evolution of a state vector or unitary matrix for a
    given Hamiltonian.

    Evolve the state vector (``psi0``) using a given Hamiltonian (``H``), by
    integrating the set of ordinary differential equations that define the
    system. Alternatively evolve a unitary matrix in solving the Schrodinger
    operator equation.

    The output is either the state vector or unitary matrix at arbitrary points
    in time (``tlist``), or the expectation values of the supplied operators
    (``e_ops``). If ``e_ops`` is a callback function, it is invoked for each
    time in ``tlist`` with time and the state as arguments, and the function
    does not use any return values. ``e_ops`` cannot be used in conjunction
    with solving the Schrodinger operator equation

    Parameters
    ----------

    H : :class:`~qutip.Qobj`, :class:`~qutip.QobjEvo`, list, or callable
        System Hamiltonian as a :obj:`~qutip.Qobj` , list of
        :obj:`~qutip.Qobj` and coefficient, :obj:`~qutip.QobjEvo`,
        or a callback function for time-dependent Hamiltonians. List format
        and options can be found in QobjEvo's description.

    psi0 : :class:`~qutip.Qobj`
        Initial state vector (ket) or initial unitary operator ``psi0 = U``.

    tlist : array_like of float
        List of times for :math:`t`.

    e_ops : None / list / callback function, optional
        A list of operators as `Qobj` and/or callable functions (can be mixed)
        or a single callable function. For callable functions, they are called
        as ``f(t, state)`` and return the expectation value. A single
        callback's expectation value can be any type, but a callback part of a
        list must return a number as the expectation value. For operators, the
        result's expect will be computed by :func:`qutip.expect` when the state
        is a ``ket``. For operator evolution, the overlap is computed by: ::

            (e_ops[i].dag() * op(t)).tr()

    args : dict, optional
        Dictionary of scope parameters for time-dependent Hamiltonians.

    options : :obj:`~qutip.solver.Options`, optional
        Options for the ODE solver.

    progress_bar : :obj:`~BaseProgressBar`, optional
        Optional instance of :obj:`~BaseProgressBar`, or a subclass thereof,
        for showing the progress of the simulation.

    Returns
    -------

    output: :class:`~qutip.solver.Result`
        An instance of the class :class:`~qutip.solver.Options`, which
        contains either an array of expectation values for the times
        specified by ``tlist``, or an array or state vectors
        corresponding to the times in ``tlist`` (if ``e_ops`` is an empty
        list), or nothing if a callback function was given inplace of
        operators for which to calculate the expectation values.
    """
    if e_ops is None:
        e_ops = []
    if isinstance(e_ops, Qobj):
        e_ops = [e_ops]
    elif isinstance(e_ops, dict):
        e_ops_dict = e_ops
        e_ops = [e for e in e_ops.values()]
    else:
        e_ops_dict = None

    if progress_bar is None:
        progress_bar = BaseProgressBar()
    if progress_bar is True:
        progress_bar = TextProgressBar()

    if not (psi0.isket or psi0.isunitary):
        raise TypeError("The unitary solver requires psi0 to be"
                        " a ket as initial state"
                        " or a unitary as initial operator.")

    if options is None:
        options = Options()
    if options.rhs_reuse and not isinstance(H, SolverSystem):
        # TODO: deprecate when going to class based solver.
        if "sesolve" in solver_safe:
            H = solver_safe["sesolve"]

    if args is None:
        args = {}

    check_use_openmp(options)

    if isinstance(H, SolverSystem):
        ss = H
    elif isinstance(H, (list, Qobj, QobjEvo)):
        ss = _sesolve_QobjEvo(H, tlist, args, options)
    elif callable(H):
        ss = _sesolve_func_td(H, args, options)
    else:
        raise Exception("Invalid H type")

    func, ode_args = ss.makefunc(ss, psi0, args, e_ops, options)

    if _safe_mode:
        v = psi0.full().ravel('F')
        func(0., v, *ode_args) + v

    res = _generic_ode_solve(func, ode_args, psi0, tlist, e_ops, options,
                             progress_bar, dims=psi0.dims)
    if e_ops_dict:
        res.expect = {e: res.expect[n]
                      for n, e in enumerate(e_ops_dict.keys())}
    return res


# -----------------------------------------------------------------------------
# A time-dependent unitary wavefunction equation on the list-function format
#
def _sesolve_QobjEvo(H, tlist, args, opt):
    """<commit_msg>fix exception type in sesolve<commit_after>"""
This module provides solvers for the unitary Schrodinger equation.
"""

__all__ = ['sesolve']

import os
import types
import numpy as np
import scipy.integrate
from scipy.linalg import norm as la_norm
from qutip.cy.stochastic import normalize_inplace
import qutip.settings as qset
from qutip.qobj import Qobj
from qutip.qobjevo import QobjEvo
from qutip.cy.spconvert import dense1D_to_fastcsr_ket, dense2D_to_fastcsr_fmode
from qutip.cy.spmatfuncs import (cy_expect_psi, cy_ode_psi_func_td,
                                cy_ode_psi_func_td_with_state)
from qutip.solver import Result, Options, config, solver_safe, SolverSystem
from qutip.superoperator import vec2mat
from qutip.ui.progressbar import (BaseProgressBar, TextProgressBar)
from qutip.cy.openmp.utilities import check_use_openmp, openmp_components

def sesolve(H, psi0, tlist, e_ops=None, args=None, options=None,
            progress_bar=None, _safe_mode=True):
    """
    Schrödinger equation evolution of a state vector or unitary matrix for a
    given Hamiltonian.

    Evolve the state vector (``psi0``) using a given Hamiltonian (``H``), by
    integrating the set of ordinary differential equations that define the
    system. Alternatively evolve a unitary matrix in solving the Schrodinger
    operator equation.

    The output is either the state vector or unitary matrix at arbitrary points
    in time (``tlist``), or the expectation values of the supplied operators
    (``e_ops``). If ``e_ops`` is a callback function, it is invoked for each
    time in ``tlist`` with time and the state as arguments, and the function
    does not use any return values. ``e_ops`` cannot be used in conjunction
    with solving the Schrodinger operator equation

    Parameters
    ----------

    H : :class:`~qutip.Qobj`, :class:`~qutip.QobjEvo`, list, or callable
        System Hamiltonian as a :obj:`~qutip.Qobj` , list of
        :obj:`~qutip.Qobj` and coefficient, :obj:`~qutip.QobjEvo`,
        or a callback function for time-dependent Hamiltonians. List format
        and options can be found in QobjEvo's description.

    psi0 : :class:`~qutip.Qobj`
        Initial state vector (ket) or initial unitary operator ``psi0 = U``.

    tlist : array_like of float
        List of times for :math:`t`.

    e_ops : None / list / callback function, optional
        A list of operators as `Qobj` and/or callable functions (can be mixed)
        or a single callable function. For callable functions, they are called
        as ``f(t, state)`` and return the expectation value. A single
        callback's expectation value can be any type, but a callback part of a
        list must return a number as the expectation value. For operators, the
        result's expect will be computed by :func:`qutip.expect` when the state
        is a ``ket``. For operator evolution, the overlap is computed by: ::

            (e_ops[i].dag() * op(t)).tr()

    args : dict, optional
        Dictionary of scope parameters for time-dependent Hamiltonians.

    options : :obj:`~qutip.solver.Options`, optional
        Options for the ODE solver.

    progress_bar : :obj:`~BaseProgressBar`, optional
        Optional instance of :obj:`~BaseProgressBar`, or a subclass thereof,
        for showing the progress of the simulation.

    Returns
    -------

    output: :class:`~qutip.solver.Result`
        An instance of the class :class:`~qutip.solver.Options`, which
        contains either an array of expectation values for the times
        specified by ``tlist``, or an array or state vectors
        corresponding to the times in ``tlist`` (if ``e_ops`` is an empty
        list), or nothing if a callback function was given inplace of
        operators for which to calculate the expectation values.
    """
    if e_ops is None:
        e_ops = []
    if isinstance(e_ops, Qobj):
        e_ops = [e_ops]
    elif isinstance(e_ops, dict):
        e_ops_dict = e_ops
        e_ops = [e for e in e_ops.values()]
    else:
        e_ops_dict = None

    if progress_bar is None:
        progress_bar = BaseProgressBar()
    if progress_bar is True:
        progress_bar = TextProgressBar()

    if not (psi0.isket or psi0.isunitary):
        raise TypeError("The unitary solver requires psi0 to be"
                        " a ket as initial state"
                        " or a unitary as initial operator.")

    if options is None:
        options = Options()
    if options.rhs_reuse and not isinstance(H, SolverSystem):
        # TODO: deprecate when going to class based solver.
        if "sesolve" in solver_safe:
            H = solver_safe["sesolve"]

    if args is None:
        args = {}

    check_use_openmp(options)

    if isinstance(H, SolverSystem):
        ss = H
    elif isinstance(H, (list, Qobj, QobjEvo)):
        ss = _sesolve_QobjEvo(H, tlist, args, options)
    elif callable(H):
        ss = _sesolve_func_td(H, args, options)
    else:
        raise TypeError(f"Invalid H type {type(H)}")

    func, ode_args = ss.makefunc(ss, psi0, args, e_ops, options)

    if _safe_mode:
        v = psi0.full().ravel('F')
        func(0., v, *ode_args) + v

    res = _generic_ode_solve(func, ode_args, psi0, tlist, e_ops, options,
                             progress_bar, dims=psi0.dims)
    if e_ops_dict:
        res.expect = {e: res.expect[n]
                      for n, e in enumerate(e_ops_dict.keys())}
    return res


# -----------------------------------------------------------------------------
# A time-dependent unitary wavefunction equation on the list-function format
#
def _sesolve_QobjEvo(H, tlist, args, opt):
    """
{'fraction_non_alphanumeric': 0.09797506410823238, 'fraction_numerical': 0.0026527544433636928, 'mean_word_length': 2.6744639376218324, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3778197', 'n_tokens_mistral': 3404, 'n_tokens_neox': 3285, 'n_words': 1219}