<commit_before>"""
Tests for qutip.nonmarkov.bofin_solvers.
"""

import numpy as np
import pytest
from numpy.linalg import eigvalsh
from scipy.integrate import quad
from scipy.sparse import csr_matrix

from qutip import (
    basis, destroy, expect, liouvillian, sigmax, sigmaz,
    tensor, Qobj, QobjEvo, Options,
)
from qutip.fastsparse import fast_csr_matrix
from qutip.nonmarkov.bofin_baths import (
    BathExponent,
    Bath,
    BosonicBath,
    DrudeLorentzBath,
    DrudeLorentzPadeBath,
    UnderDampedBath,
    FermionicBath,
    LorentzianBath,
    LorentzianPadeBath,
)
from qutip.nonmarkov.bofin_solvers import (
    HierarchyADOs,
    HierarchyADOsState,
    HEOMSolver,
    HSolverDL,
    _GatherHEOMRHS,
)
from qutip.ui.progressbar import BaseProgressBar, TextProgressBar


class TestHierarchyADOs:
    def mk_exponents(self, dims):
        return [
            BathExponent("I", dim, Q=None, ck=1.0, vk=2.0) for dim in dims
        ]

    def test_create(self):
        exponents = self.mk_exponents([2, 3])
        ados = HierarchyADOs(exponents, max_depth=2)

        assert ados.exponents == exponents
        assert ados.max_depth == 2

        assert ados.dims == [2, 3]
        assert ados.vk == [2.0, 2.0]
        assert ados.ck == [1.0, 1.0]
        assert ados.ck2 == [None, None]
        assert ados.sigma_bar_k_offset == [None, None]

        assert ados.labels == [
            (0, 0), (0, 1), (0, 2), (1, 0), (1, 1),
        ]

    def test_state_idx(self):
        ados = HierarchyADOs(self.mk_exponents([2, 3]), max_depth=2)
        assert ados.idx((0, 0)) == 0
        assert ados.idx((0, 1)) == 1
        assert ados.idx((0, 2)) == 2
        assert ados.idx((1, 0)) == 3
        assert ados.idx((1, 1)) == 4

    def test_next(self):
        ados = HierarchyADOs(self.mk_exponents([2, 3]), max_depth=2)
        assert ados.next((0, 0), 0) == (1, 0)
        assert ados.next((0, 0), 1) == (0, 1)
        assert ados.next((1, 0), 0) is None
        assert ados.next((1, 0), 1) == (1, 1)
        assert ados.next((1, 1), 1) is None

    def test_prev(self):
        ados = HierarchyADOs(self.mk_exponents([2, 3]), max_depth=2)
        assert ados.prev((0, 0), 0) is None
        assert ados.prev((0, 0), 1) is None
        assert ados.prev((1, 0), 0) == (0, 0)
        assert ados.prev((0, 1), 1) == (0, 0)
        assert ados.prev((1, 1), 1) == (1, 0)
        assert ados.prev((0, 2), 1) == (0, 1)

    def test_exps(self):
        ados = HierarchyADOs(self.mk_exponents([3, 3, 2]), max_depth=4)
        assert ados.exps((0, 0, 0)) == ()
        assert ados.exps((1, 0, 0)) == (ados.exponents[0],)
        assert ados.exps((2, 0, 0)) == (
            ados.exponents[0], ados.exponents[0],
        )
        assert ados.exps((1, 2, 1)) == (
            ados.exponents[0],
            ados.exponents[1], ados.exponents[1],
            ados.exponents[2],
        )

    def test_filter_by_nothing(self):
        ados = HierarchyADOs(self.mk_exponents([2, 3]), max_depth=2)
        assert ados.filter() == [
            (0, 0), (0, 1), (0, 2), (1, 0), (1, 1),
        ]

    def test_filter_by_level(self):
        ados = HierarchyADOs(self.mk_exponents([2, 3]), max_depth=2)
        assert ados.filter(level=0) == [
            (0, 0),
        ]
        assert ados.filter(level=1) == [
            (0, 1),
            (1, 0),
        ]
        assert ados.filter(level=2) == [
            (0, 2),
            (1, 1),
        ]
        assert ados.filter(level=3) == []

    def test_filter_by_exponents(self):
        ados = HierarchyADOs(self.mk_exponents([2, 3]), max_depth=2)
        assert ados.filter(dims=[]) == [
            (0, 0),
        ]
        assert ados.filter(dims=[2]) == [
            (1, 0),
        ]
        assert ados.filter(level=1, dims=[2]) == [
            (1, 0),
        ]
        assert ados.filter(dims=[3]) == [
            (0, 1),
        ]
        assert ados.filter(dims=[2, 3]) == [
            (1, 1),
        ]
        assert ados.filter(level=2, dims=[2, 3]) == [
            (1, 1),
        ]
        assert ados.filter(dims=[3, 3]) == [
            (0, 2),
        ]
        assert ados.filter(types=["I"]) == [
            (0, 1),
            (1, 0),
        ]
        assert ados.filter(types=["I", "I"]) == [
            (0, 2),
            (1, 1),
        ]

        with pytest.raises(ValueError) as err:
            ados.filter(types=[], dims=[2])
        assert str(err.value) == (
            "The tags, dims and types filters must all be the same length."
        )

        with pytest.raises(ValueError) as err:
            ados.filter(dims=[2, 2, 2])
        assert str(err.value) == (
            "The maximum depth for the hierarchy is 2 but 3 levels of"
            " excitation filters were given."
        )

        with pytest.raises(ValueError) as err:
            ados.filter(level=0, dims=[2])
        assert str(err.value) == (
            "The level parameter is 0 but 1 levels of excitation filters"
            " were given."
        )


class TestHierarchyADOsState:
    def mk_ados(self, bath_dims, max_depth):
        exponents = [
            BathExponent("I", dim, Q=None, ck=1.0, vk=2.0) for dim in bath_dims
        ]
        ados = HierarchyADOs(exponents, max_depth=max_depth)
        return ados

    def mk_rho_and_soln(self, ados, rho_dims):
        n_ados = len(ados.labels)
        ado_soln = np.random.rand(n_ados, *[2**d for d in rho_dims])
        rho = Qobj(ado_soln[0, :], dims=[2, 2])
        return rho, ado_soln

    def test_create(self):
        ados = self.mk_ados([2, 3], max_depth=2)
        rho, ado_soln = self.mk_rho_and_soln(ados, [2, 2])
        ado_state = HierarchyADOsState(rho, ados, ado_soln)
        assert ado_state.rho == rho
        assert ado_state.labels == ados.labels
        assert ado_state.exponents == ados.exponents
        assert ado_state.idx((0, 0)) == ados.idx((0, 0))
        assert ado_state.idx((0, 1)) == ados.idx((0, 1))

    def test_extract(self):
        ados = self.mk_ados([2, 3], max_depth=2)
        rho, ado_soln = self.mk_rho_and_soln(ados, [2, 2])
        ado_state = HierarchyADOsState(rho, ados, ado_soln)
        ado_state.extract((0, 0)) == rho
        ado_state.extract(0) == rho
        ado_state.extract((0, 1)) == Qobj(ado_soln[1, :], dims=rho.dims)
        ado_state.extract(1) == Qobj(ado_soln[1, :], dims=rho.dims)


class DrudeLorentzPureDephasingModel:
    """ Analytic Drude-Lorentz pure-dephasing model for testing the HEOM solver.
    """
    def __init__(self, lam, gamma, T, Nk):
        self.lam = lam
        self.gamma = gamma
        self.T = T
        self.Nk = Nk
        # we add a very weak system hamiltonian here to avoid having
        # singular system that causes problems for the scipy.sparse.linalg
        # superLU solver used in spsolve.
        self.H = Qobj(1e-5 * np.ones((2, 2)))
        self.Q = sigmaz()

    def rho(self):
        """ Initial state. """
        return 0.5 * Qobj(np.ones((2, 2)))

    def state_results(self, states):
        projector = basis(2, 0) * basis(2, 1).dag()
        return expect(states, projector)

    def analytic_results(self, tlist):
        lam, gamma, T = self.lam, self.gamma, self.T<commit_msg>Correctly set density matrix dimensions when constructing ADOs in HEOM ADO state tests.<commit_after>"""
Tests for qutip.nonmarkov.bofin_solvers.
"""

import numpy as np
import pytest
from numpy.linalg import eigvalsh
from scipy.integrate import quad
from scipy.sparse import csr_matrix

from qutip import (
    basis, destroy, expect, liouvillian, sigmax, sigmaz,
    tensor, Qobj, QobjEvo, Options,
)
from qutip.fastsparse import fast_csr_matrix
from qutip.nonmarkov.bofin_baths import (
    BathExponent,
    Bath,
    BosonicBath,
    DrudeLorentzBath,
    DrudeLorentzPadeBath,
    UnderDampedBath,
    FermionicBath,
    LorentzianBath,
    LorentzianPadeBath,
)
from qutip.nonmarkov.bofin_solvers import (
    HierarchyADOs,
    HierarchyADOsState,
    HEOMSolver,
    HSolverDL,
    _GatherHEOMRHS,
)
from qutip.ui.progressbar import BaseProgressBar, TextProgressBar


class TestHierarchyADOs:
    def mk_exponents(self, dims):
        return [
            BathExponent("I", dim, Q=None, ck=1.0, vk=2.0) for dim in dims
        ]

    def test_create(self):
        exponents = self.mk_exponents([2, 3])
        ados = HierarchyADOs(exponents, max_depth=2)

        assert ados.exponents == exponents
        assert ados.max_depth == 2

        assert ados.dims == [2, 3]
        assert ados.vk == [2.0, 2.0]
        assert ados.ck == [1.0, 1.0]
        assert ados.ck2 == [None, None]
        assert ados.sigma_bar_k_offset == [None, None]

        assert ados.labels == [
            (0, 0), (0, 1), (0, 2), (1, 0), (1, 1),
        ]

    def test_state_idx(self):
        ados = HierarchyADOs(self.mk_exponents([2, 3]), max_depth=2)
        assert ados.idx((0, 0)) == 0
        assert ados.idx((0, 1)) == 1
        assert ados.idx((0, 2)) == 2
        assert ados.idx((1, 0)) == 3
        assert ados.idx((1, 1)) == 4

    def test_next(self):
        ados = HierarchyADOs(self.mk_exponents([2, 3]), max_depth=2)
        assert ados.next((0, 0), 0) == (1, 0)
        assert ados.next((0, 0), 1) == (0, 1)
        assert ados.next((1, 0), 0) is None
        assert ados.next((1, 0), 1) == (1, 1)
        assert ados.next((1, 1), 1) is None

    def test_prev(self):
        ados = HierarchyADOs(self.mk_exponents([2, 3]), max_depth=2)
        assert ados.prev((0, 0), 0) is None
        assert ados.prev((0, 0), 1) is None
        assert ados.prev((1, 0), 0) == (0, 0)
        assert ados.prev((0, 1), 1) == (0, 0)
        assert ados.prev((1, 1), 1) == (1, 0)
        assert ados.prev((0, 2), 1) == (0, 1)

    def test_exps(self):
        ados = HierarchyADOs(self.mk_exponents([3, 3, 2]), max_depth=4)
        assert ados.exps((0, 0, 0)) == ()
        assert ados.exps((1, 0, 0)) == (ados.exponents[0],)
        assert ados.exps((2, 0, 0)) == (
            ados.exponents[0], ados.exponents[0],
        )
        assert ados.exps((1, 2, 1)) == (
            ados.exponents[0],
            ados.exponents[1], ados.exponents[1],
            ados.exponents[2],
        )

    def test_filter_by_nothing(self):
        ados = HierarchyADOs(self.mk_exponents([2, 3]), max_depth=2)
        assert ados.filter() == [
            (0, 0), (0, 1), (0, 2), (1, 0), (1, 1),
        ]

    def test_filter_by_level(self):
        ados = HierarchyADOs(self.mk_exponents([2, 3]), max_depth=2)
        assert ados.filter(level=0) == [
            (0, 0),
        ]
        assert ados.filter(level=1) == [
            (0, 1),
            (1, 0),
        ]
        assert ados.filter(level=2) == [
            (0, 2),
            (1, 1),
        ]
        assert ados.filter(level=3) == []

    def test_filter_by_exponents(self):
        ados = HierarchyADOs(self.mk_exponents([2, 3]), max_depth=2)
        assert ados.filter(dims=[]) == [
            (0, 0),
        ]
        assert ados.filter(dims=[2]) == [
            (1, 0),
        ]
        assert ados.filter(level=1, dims=[2]) == [
            (1, 0),
        ]
        assert ados.filter(dims=[3]) == [
            (0, 1),
        ]
        assert ados.filter(dims=[2, 3]) == [
            (1, 1),
        ]
        assert ados.filter(level=2, dims=[2, 3]) == [
            (1, 1),
        ]
        assert ados.filter(dims=[3, 3]) == [
            (0, 2),
        ]
        assert ados.filter(types=["I"]) == [
            (0, 1),
            (1, 0),
        ]
        assert ados.filter(types=["I", "I"]) == [
            (0, 2),
            (1, 1),
        ]

        with pytest.raises(ValueError) as err:
            ados.filter(types=[], dims=[2])
        assert str(err.value) == (
            "The tags, dims and types filters must all be the same length."
        )

        with pytest.raises(ValueError) as err:
            ados.filter(dims=[2, 2, 2])
        assert str(err.value) == (
            "The maximum depth for the hierarchy is 2 but 3 levels of"
            " excitation filters were given."
        )

        with pytest.raises(ValueError) as err:
            ados.filter(level=0, dims=[2])
        assert str(err.value) == (
            "The level parameter is 0 but 1 levels of excitation filters"
            " were given."
        )


class TestHierarchyADOsState:
    def mk_ados(self, bath_dims, max_depth):
        exponents = [
            BathExponent("I", dim, Q=None, ck=1.0, vk=2.0) for dim in bath_dims
        ]
        ados = HierarchyADOs(exponents, max_depth=max_depth)
        return ados

    def mk_rho_and_soln(self, ados, rho_dims):
        n_ados = len(ados.labels)
        ado_soln = np.random.rand(n_ados, *[np.product(d) for d in rho_dims])
        rho = Qobj(ado_soln[0, :], dims=rho_dims)
        return rho, ado_soln

    def test_create(self):
        ados = self.mk_ados([2, 3], max_depth=2)
        rho, ado_soln = self.mk_rho_and_soln(ados, [[2], [2]])
        ado_state = HierarchyADOsState(rho, ados, ado_soln)
        assert ado_state.rho == rho
        assert ado_state.labels == ados.labels
        assert ado_state.exponents == ados.exponents
        assert ado_state.idx((0, 0)) == ados.idx((0, 0))
        assert ado_state.idx((0, 1)) == ados.idx((0, 1))

    def test_extract(self):
        ados = self.mk_ados([2, 3], max_depth=2)
        rho, ado_soln = self.mk_rho_and_soln(ados, [[2], [2]])
        ado_state = HierarchyADOsState(rho, ados, ado_soln)
        ado_state.extract((0, 0)) == rho
        ado_state.extract(0) == rho
        ado_state.extract((0, 1)) == Qobj(ado_soln[1, :], dims=rho.dims)
        ado_state.extract(1) == Qobj(ado_soln[1, :], dims=rho.dims)


class DrudeLorentzPureDephasingModel:
    """ Analytic Drude-Lorentz pure-dephasing model for testing the HEOM solver.
    """
    def __init__(self, lam, gamma, T, Nk):
        self.lam = lam
        self.gamma = gamma
        self.T = T
        self.Nk = Nk
        # we add a very weak system hamiltonian here to avoid having
        # singular system that causes problems for the scipy.sparse.linalg
        # superLU solver used in spsolve.
        self.H = Qobj(1e-5 * np.ones((2, 2)))
        self.Q = sigmaz()

    def rho(self):
        """ Initial state. """
        return 0.5 * Qobj(np.ones((2, 2)))

    def state_results(self, states):
        projector = basis(2, 0) * basis(2, 1).dag()
        return expect(states, projector)

    def analytic_results(self, tlist):
        lam, gamma, T = self.lam, self.gamma, self.T
{'fraction_non_alphanumeric': 0.14699224912545442, 'fraction_numerical': 0.033404211537142466, 'mean_word_length': 2.0902924968206866, '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': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5429474', 'n_tokens_mistral': 5614, 'n_tokens_neox': 5278, 'n_words': 1218}