<commit_before># This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, <NAME> and <NAME>.
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################

import pytest
import functools
import numpy as np
import qutip


def _return_constant(constant, *args, **kwargs):
    return constant


def _return_decay(constant, rate, t, *args, **kwargs):
    return constant * np.exp(-rate * t)


@pytest.mark.usefixtures("in_temporary_directory")
class StatesAndExpectOutputCase:
    """
    Mixin class to test the states and expectation values from ``mcsolve``.
    """
    size = 10
    h = qutip.num(size)
    state = qutip.basis(size, size-1)
    times = np.linspace(0, 1, 101)
    e_ops = [qutip.num(size)]
    ntraj = 750

    def _assert_states(self, result, expected, tol):
        assert hasattr(result, 'states')
        assert len(result.states) == len(self.times)
        for test_operator, expected_part in zip(self.e_ops, expected):
            test = qutip.expect(test_operator, result.states)
            np.testing.assert_allclose(test, expected_part, rtol=tol)

    def _assert_expect(self, result, expected, tol):
        assert hasattr(result, 'expect')
        assert len(result.expect) == len(self.e_ops)
        for test, expected_part in zip(result.expect, expected):
            np.testing.assert_allclose(test, expected_part, rtol=tol)

    def test_states(self, hamiltonian, c_ops, expected, tol):
        options = qutip.Options(average_states=True, store_states=True)
        result = qutip.mcsolve(hamiltonian, self.state, self.times,
                               c_ops=c_ops, e_ops=[], ntraj=self.ntraj,
                               options=options)
        self._assert_states(result, expected, tol)

    def test_expect(self, hamiltonian, c_ops, expected, tol):
        result = qutip.mcsolve(hamiltonian, self.state, self.times,
                               c_ops=c_ops, e_ops=self.e_ops, ntraj=self.ntraj)
        self._assert_expect(result, expected, tol)

    def test_states_and_expect(self, hamiltonian, c_ops, expected, tol):
        options = qutip.Options(average_states=True, store_states=True)
        result = qutip.mcsolve(hamiltonian, self.state, self.times,
                               c_ops=c_ops, e_ops=self.e_ops, ntraj=self.ntraj,
                               options=options)
        self._assert_expect(result, expected, tol)
        self._assert_states(result, expected, tol)


class TestNoCollapse(StatesAndExpectOutputCase):
    """
    Test that `mcsolve` correctly solves the system when there is a constant
    Hamiltonian and no collapses.
    """
    def pytest_generate_tests(self, metafunc):
        tol = 1e-8
        expect = (qutip.expect(self.e_ops[0], self.state)
                  * np.ones_like(self.times))
        # Use partial to allow pickling.
        zero_function = functools.partial(_return_constant, 0)
        hamiltonian_types = [
            (self.h, "Qobj", False),
            ([self.h], "list", False),
            ([self.h, [self.h, '0']], "string", True),
            ([self.h, [self.h, zero_function]], "function", False),
        ]
        cases = []
        for hamiltonian, id, slow in hamiltonian_types:
            marks = [pytest.mark.slow] if slow else []
            cases.append(pytest.param(hamiltonian, [], [expect], tol,
                                      id=id, marks=marks))
        metafunc.parametrize(['hamiltonian', 'c_ops', 'expected', 'tol'],
                             cases)


class TestConstantCollapse(StatesAndExpectOutputCase):
    """
    Test that `mcsolve` correctly solves the system when there is a constant
    collapse operator.
    """
    def pytest_generate_tests(self, metafunc):
        tol = 0.05
        coupling = 0.2
        expect = (qutip.expect(self.e_ops[0], self.state)
                  * np.exp(-coupling * self.times))
        collapse_op = qutip.destroy(self.size)
        # Use partial to allow pickling.
        collapse_function = functools.partial(_return_constant,
                                              np.sqrt(coupling))
        c_op_types = [<commit_msg>Reduce number of tests in test_mcsolve.py<commit_after># This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, <NAME> and <NAME>.
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################

import pytest
import functools
import numpy as np
import qutip


def _return_constant(constant, *args, **kwargs):
    return constant


def _return_decay(constant, rate, t, *args, **kwargs):
    return constant * np.exp(-rate * t)


@pytest.mark.usefixtures("in_temporary_directory")
class StatesAndExpectOutputCase:
    """
    Mixin class to test the states and expectation values from ``mcsolve``.
    """
    size = 10
    h = qutip.num(size)
    state = qutip.basis(size, size-1)
    times = np.linspace(0, 1, 101)
    e_ops = [qutip.num(size)]
    ntraj = 750

    def _assert_states(self, result, expected, tol):
        assert hasattr(result, 'states')
        assert len(result.states) == len(self.times)
        for test_operator, expected_part in zip(self.e_ops, expected):
            test = qutip.expect(test_operator, result.states)
            np.testing.assert_allclose(test, expected_part, rtol=tol)

    def _assert_expect(self, result, expected, tol):
        assert hasattr(result, 'expect')
        assert len(result.expect) == len(self.e_ops)
        for test, expected_part in zip(result.expect, expected):
            np.testing.assert_allclose(test, expected_part, rtol=tol)

    def test_states_and_expect(self, hamiltonian, c_ops, expected, tol):
        options = qutip.Options(average_states=True, store_states=True)
        result = qutip.mcsolve(hamiltonian, self.state, self.times,
                               c_ops=c_ops, e_ops=self.e_ops, ntraj=self.ntraj,
                               options=options)
        self._assert_expect(result, expected, tol)
        self._assert_states(result, expected, tol)


class TestNoCollapse(StatesAndExpectOutputCase):
    """
    Test that `mcsolve` correctly solves the system when there is a constant
    Hamiltonian and no collapses.
    """
    def pytest_generate_tests(self, metafunc):
        tol = 1e-8
        expect = (qutip.expect(self.e_ops[0], self.state)
                  * np.ones_like(self.times))
        # Use partial to allow pickling.
        zero_function = functools.partial(_return_constant, 0)
        hamiltonian_types = [
            (self.h, "Qobj", False),
            ([self.h], "list", False),
            ([self.h, [self.h, '0']], "string", True),
            ([self.h, [self.h, zero_function]], "function", False),
        ]
        cases = []
        for hamiltonian, id, slow in hamiltonian_types:
            if slow and 'only' in metafunc.function.__name__:
                # Skip the single-output test if it's a slow case.
                continue
            marks = [pytest.mark.slow] if slow else []
            cases.append(pytest.param(hamiltonian, [], [expect], tol,
                                      id=id, marks=marks))
        metafunc.parametrize(['hamiltonian', 'c_ops', 'expected', 'tol'],
                             cases)

    # Previously the "states_only" and "expect_only" tests were mixed in to
    # every other test case.  We move them out into the simplest set so that
    # their behaviour remains tested, but isn't repeated as often to keep test
    # runtimes shorter.  The known-good cases are still tested in the other
    # test cases, this is just testing the single-output behaviour.

    def test_states_only(self, hamiltonian, c_ops, expected, tol):
        options = qutip.Options(average_states=True, store_states=True)
        result = qutip.mcsolve(hamiltonian, self.state, self.times,
                               c_ops=c_ops, e_ops=[], ntraj=self.ntraj,
                               options=options)
        self._assert_states(result, expected, tol)

    def test_expect_only(self, hamiltonian, c_ops, expected, tol):
        result = qutip.mcsolve(hamiltonian, self.state, self.times,
                               c_ops=c_ops, e_ops=self.e_ops, ntraj=self.ntraj)
        self._assert_expect(result, expected, tol)


class TestConstantCollapse(StatesAndExpectOutputCase):
    """
    Test that `mcsolve` correctly solves the system when there is a constant
    collapse operator.
    """
    def pytest_generate_tests(self, metafunc):
        tol = 0.05
        coupling = 0.2
        expect = (qutip.expect(self.e_ops[0], self.state)
                  * np.exp(-coupling * self.times))
        collapse_op = qutip.destroy(self.size)
        # Use partial to allow pickling.
        collapse_function = functools.partial(_return_constant,
                                              np.sqrt(coupling))
        c_op_types = [
{'fraction_non_alphanumeric': 0.10193974411886092, 'fraction_numerical': 0.004787453569954601, 'mean_word_length': 2.7360468701819305, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6315588', 'n_tokens_mistral': 3608, 'n_tokens_neox': 3420, 'n_words': 1099}