def test_discrete_level_model_fermionic_bath(self, evo, liouvillianize):
        dlm = DiscreteLevelCurrentModel(
            gamma=0.01,
            W=1,
            T=0.025851991,
            lmax=10,
        )
        H_sys = hamiltonian_to_sys(dlm.H, evo, liouvillianize)
        ck_plus, vk_plus, ck_minus, vk_minus = dlm.bath_coefficients()

        options = Options(
            nsteps=15_000,
            store_states=True,
            rtol=1e-7,
            atol=1e-7,
        )
        bath = FermionicBath(dlm.Q, ck_plus, vk_plus, ck_minus, vk_minus)
        # for a single impurity we converge with max_depth = 2
        hsolver = HEOMSolver(H_sys, bath, 2, options=options)

        tlist = [0, 600]
        result = hsolver.run(dlm.rho(), tlist, ado_return=True)
        current = dlm.state_current(result.ado_states[-1])
        analytic_current = dlm.analytic_current()
        np.testing.assert_allclose(analytic_current, current, rtol=1e-3)

        rho_final, ado_state = hsolver.steady_state()
        current = dlm.state_current(ado_state)
        analytic_current = dlm.analytic_current()
        np.testing.assert_allclose(analytic_current, current, rtol=1e-3)
    def test_underdamped_pure_dephasing_model_underdamped_bath(
            self, atol=1e-3):
        udm = UnderdampedPureDephasingModel(
            lam=0.1,
            gamma=0.05,
            w0=1,
            T=1 / 0.95,
            Nk=2,
        )
        bath = UnderDampedBath(
            Q=udm.Q,
            lam=udm.lam,
            T=udm.T,
            Nk=udm.Nk,
            gamma=udm.gamma,
            w0=udm.w0,
        )

        options = Options(nsteps=15000, store_states=True)
        hsolver = HEOMSolver(udm.H, bath, 14, options=options)

        tlist = np.linspace(0, 10, 21)
        result = hsolver.run(udm.rho(), tlist)

        test = udm.state_results(result.states)
        expected = udm.analytic_results(tlist)
        np.testing.assert_allclose(test, expected, atol=atol)

        rho_final, ado_state = hsolver.steady_state()
        test = udm.state_results([rho_final])
        expected = udm.analytic_results([5000])
        np.testing.assert_allclose(test, expected, atol=atol)
        assert rho_final == ado_state.extract(0)
    def test_pure_dephasing_model_bosonic_bath(self,
                                               evo,
                                               combine,
                                               liouvillianize,
                                               atol=1e-3):
        dlm = DrudeLorentzPureDephasingModel(
            lam=0.025,
            gamma=0.05,
            T=1 / 0.95,
            Nk=2,
        )
        ck_real, vk_real, ck_imag, vk_imag = dlm.bath_coefficients()
        H_sys = hamiltonian_to_sys(dlm.H, evo, liouvillianize)

        bath = BosonicBath(dlm.Q, ck_real, vk_real, ck_imag, vk_imag)
        options = Options(nsteps=15000, store_states=True)
        hsolver = HEOMSolver(H_sys, bath, 14, options=options)

        tlist = np.linspace(0, 10, 21)
        result = hsolver.run(dlm.rho(), tlist)

        test = dlm.state_results(result.states)
        expected = dlm.analytic_results(tlist)
        np.testing.assert_allclose(test, expected, atol=atol)

        rho_final, ado_state = hsolver.steady_state()
        test = dlm.state_results([rho_final])
        expected = dlm.analytic_results([100])
        np.testing.assert_allclose(test, expected, atol=atol)
        assert rho_final == ado_state.extract(0)
    def test_integration_error(self):
        dlm = DrudeLorentzPureDephasingModel(
            lam=0.025,
            gamma=0.05,
            T=1 / 0.95,
            Nk=2,
        )
        ck_real, vk_real, ck_imag, vk_imag = dlm.bath_coefficients()

        bath = BosonicBath(dlm.Q, ck_real, vk_real, ck_imag, vk_imag)
        options = Options(nsteps=10)
        hsolver = HEOMSolver(dlm.H, bath, 14, options=options)

        with pytest.raises(RuntimeError) as err:
            hsolver.run(dlm.rho(), tlist=[0, 10])

        assert str(err.value) == (
            "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).")
    def test_discrete_level_model_lorentzian_baths(self,
                                                   bath_cls,
                                                   analytic_current,
                                                   atol=1e-3):
        dlm = DiscreteLevelCurrentModel(
            gamma=0.01,
            W=1,
            T=0.025851991,
            lmax=10,
        )

        options = Options(
            nsteps=15_000,
            store_states=True,
            rtol=1e-7,
            atol=1e-7,
        )
        bath_l = bath_cls(
            dlm.Q,
            gamma=dlm.gamma,
            w=dlm.W,
            T=dlm.T,
            mu=dlm.theta / 2,
            Nk=dlm.lmax,
        )
        bath_r = bath_cls(
            dlm.Q,
            gamma=dlm.gamma,
            w=dlm.W,
            T=dlm.T,
            mu=-dlm.theta / 2,
            Nk=dlm.lmax,
        )
        # for a single impurity we converge with max_depth = 2
        hsolver = HEOMSolver(dlm.H, [bath_r, bath_l], 2, options=options)

        tlist = [0, 600]
        result = hsolver.run(dlm.rho(), tlist, ado_return=True)
        current = dlm.state_current(result.ado_states[-1])
        # analytic_current = dlm.analytic_current()
        np.testing.assert_allclose(analytic_current, current, rtol=1e-3)

        rho_final, ado_state = hsolver.steady_state()
        current = dlm.state_current(ado_state)
        # analytic_current = dlm.analytic_current()
        np.testing.assert_allclose(analytic_current, current, rtol=1e-3)
    def test_pure_dephasing_model_drude_lorentz_baths(self,
                                                      terminator,
                                                      bath_cls,
                                                      atol=1e-3):
        dlm = DrudeLorentzPureDephasingModel(
            lam=0.025,
            gamma=0.05,
            T=1 / 0.95,
            Nk=2,
        )
        bath = bath_cls(
            Q=dlm.Q,
            lam=dlm.lam,
            gamma=dlm.gamma,
            T=dlm.T,
            Nk=dlm.Nk,
        )
        if terminator:
            _, terminator_op = bath.terminator()
            H_sys = liouvillian(dlm.H) + terminator_op
        else:
            H_sys = dlm.H

        options = Options(nsteps=15000, store_states=True)
        hsolver = HEOMSolver(H_sys, bath, 14, options=options)

        tlist = np.linspace(0, 10, 21)
        result = hsolver.run(dlm.rho(), tlist)

        test = dlm.state_results(result.states)
        expected = dlm.analytic_results(tlist)
        np.testing.assert_allclose(test, expected, atol=atol)

        rho_final, ado_state = hsolver.steady_state()
        test = dlm.state_results([rho_final])
        expected = dlm.analytic_results([100])
        np.testing.assert_allclose(test, expected, atol=atol)
        assert rho_final == ado_state.extract(0)
    def test_ado_return_and_ado_init(self, ado_format):
        dlm = DrudeLorentzPureDephasingModel(
            lam=0.025,
            gamma=0.05,
            T=1 / 0.95,
            Nk=2,
        )
        ck_real, vk_real, ck_imag, vk_imag = dlm.bath_coefficients()

        bath = BosonicBath(dlm.Q, ck_real, vk_real, ck_imag, vk_imag)
        options = Options(nsteps=15_000, store_states=True)
        hsolver = HEOMSolver(dlm.H, bath, 6, options=options)

        tlist_1 = [0, 1, 2]
        result_1 = hsolver.run(dlm.rho(), tlist_1, ado_return=True)

        tlist_2 = [2, 3, 4]
        rho0 = result_1.ado_states[-1]
        if ado_format == "numpy":
            rho0 = rho0._ado_state  # extract the raw numpy array
        result_2 = hsolver.run(
            rho0,
            tlist_2,
            ado_return=True,
            ado_init=True,
        )

        tlist_full = tlist_1 + tlist_2[1:]
        result_full = hsolver.run(dlm.rho(), tlist_full, ado_return=True)

        times_12 = result_1.times + result_2.times[1:]
        times_full = result_full.times
        assert times_12 == tlist_full
        assert times_full == tlist_full

        ado_states_12 = result_1.ado_states + result_2.ado_states[1:]
        ado_states_full = result_full.ado_states
        assert len(ado_states_12) == len(tlist_full)
        assert len(ado_states_full) == len(tlist_full)
        for ado_12, ado_full in zip(ado_states_12, ado_states_full):
            for label in hsolver.ados.labels:
                np.testing.assert_allclose(
                    ado_12.extract(label).full(),
                    ado_full.extract(label).full(),
                    atol=1e-6,
                )

        states_12 = result_1.states + result_2.states[1:]
        states_full = result_full.states
        assert len(states_12) == len(tlist_full)
        assert len(states_full) == len(tlist_full)
        for ado_12, state_12 in zip(ado_states_12, states_12):
            assert ado_12.rho == state_12
        for ado_full, state_full in zip(ado_states_full, states_full):
            assert ado_full.rho == state_full

        expected = dlm.analytic_results(tlist_full)
        test_12 = dlm.state_results(states_12)
        np.testing.assert_allclose(test_12, expected, atol=1e-3)
        test_full = dlm.state_results(states_full)
        np.testing.assert_allclose(test_full, expected, atol=1e-3)