support for brmesolve, there is a version of numpy that should fail, …

…but could not reproduce locally

qutip/core/environment.py

@@ -850,12 +850,14 @@ class DrudeLorentzEnvironment(BosonicEnvironment):
     """
 
     def __init__(
-        self, T: float, lam: float, gamma: float, *, tag: Any = None
+        self, T: float, lam: float, gamma: float, *, tag: Any = None,
+        Nk: int = None
     ):
         super().__init__(T, tag)
 
         self.lam = lam
         self.gamma = gamma
+        self.Nk = Nk
 
     def spectral_density(self, w: float | ArrayLike) -> (float | ArrayLike):
         """
@@ -900,6 +902,8 @@ def correlation_function(
 
         t = np.asarray(t, dtype=float)
         abs_t = np.abs(t)
+        if (self.Nk is not None):
+            Nk = self.Nk
         ck_real, vk_real, ck_imag, vk_imag = self._pade_params(Nk)
 
         def C(c, v):

qutip/solver/brmesolve.py

@@ -21,6 +21,7 @@
 from .options import _SolverOptions
 from ._feedback import _QobjFeedback, _DataFeedback
 from ..typing import EopsLike, QobjEvoLike, CoefficientLike
+from ..core.environment import BosonicEnvironment, FermionicEnvironment
 
 
 def brmesolve(
@@ -198,17 +199,23 @@ def brmesolve(
             a_op = QobjEvo(a_op, args=args, tlist=tlist)
         if isinstance(spectra, str):
             new_a_ops.append(
-                (a_op, coefficient(spectra, args={**args, 'w':0})))
+                (a_op, coefficient(spectra, args={**args, 'w': 0})))
         elif isinstance(spectra, InterCoefficient):
             new_a_ops.append((a_op, SpectraCoefficient(spectra)))
         elif isinstance(spectra, Coefficient):
             new_a_ops.append((a_op, spectra))
+        elif (isinstance(spectra, BosonicEnvironment) or
+              isinstance(spectra, FermionicEnvironment)):
+            spec = SpectraCoefficient(
+                coefficient(lambda w: spectra.power_spectrum(w))
+            )
+            new_a_ops.append((a_op, spec))
         elif callable(spectra):
             sig = inspect.signature(spectra)
             if tuple(sig.parameters.keys()) == ("w",):
                 spec = SpectraCoefficient(coefficient(spectra))
             else:
-                spec = coefficient(spectra, args={**args, 'w':0})
+                spec = coefficient(spectra, args={**args, 'w': 0})
             new_a_ops.append((a_op, spec))
         else:
             raise TypeError("a_ops's spectra not known")
@@ -274,7 +281,7 @@ class BRSolver(Solver):
     name = "brmesolve"
     solver_options = {
         "progress_bar": "",
-        "progress_kwargs": {"chunk_size":10},
+        "progress_kwargs": {"chunk_size": 10},
         "store_final_state": False,
         "store_states": None,
         "normalize_output": False,
@@ -331,7 +338,6 @@ def __init__(
         self.stats = self._initialize_stats()
         self.rhs._register_feedback({}, solver=self.name)
 
-
     def _initialize_stats(self):
         stats = super()._initialize_stats()
         stats.update({

qutip/tests/solver/test_brmesolve.py

@@ -2,6 +2,7 @@
 import numpy as np
 import qutip
 from qutip.solver.brmesolve import brmesolve
+from qutip.core.environment import DrudeLorentzEnvironment
 
 
 def pauli_spin_operators():
@@ -10,7 +11,7 @@ def pauli_spin_operators():
 
 _simple_qubit_gamma = 0.25
 coeff = qutip.coefficient(lambda t, w: _simple_qubit_gamma * (w >= 0),
-                          args={'w':0})
+                          args={'w': 0})
 _m_c_op = np.sqrt(_simple_qubit_gamma) * qutip.sigmam()
 _z_c_op = np.sqrt(_simple_qubit_gamma) * qutip.sigmaz()
 _x_a_op = [qutip.sigmax(), coeff]
@@ -132,8 +133,8 @@ def test_tensor_system():
          + w3/2. * qutip.tensor(id2, id2, qutip.sigmaz()))
 
     # White noise
-    S2 = qutip.coefficient(lambda t, w: gamma2, args={'w':0})
-    S3 = qutip.coefficient(lambda t, w: gamma3, args={'w':0})
+    S2 = qutip.coefficient(lambda t, w: gamma2, args={'w': 0})
+    S3 = qutip.coefficient(lambda t, w: gamma3, args={'w': 0})
 
     qubit_2_x = qutip.tensor(id2, qutip.sigmax(), id2)
     qubit_3_x = qutip.tensor(id2, id2, qutip.sigmax())
@@ -209,10 +210,10 @@ def _string_w_interpolating_t(kappa, times):
 
 @pytest.mark.slow
 @pytest.mark.parametrize("time_dependence_tuple", [
-        _mixed_string,
-        _separate_strings,
-        _string_w_interpolating_t,
-    ])
+    _mixed_string,
+    _separate_strings,
+    _string_w_interpolating_t,
+])
 def test_time_dependence_tuples(time_dependence_tuple):
     N = 10
     a = qutip.destroy(N)
@@ -312,3 +313,27 @@ def test_feedback():
         args={"A": qutip.BRSolver.ExpectFeedback(qutip.num(N))}
     )
     assert np.all(result.expect[0] > 4. - tol)
+
+
+@pytest.mark.parametrize("lam,gamma,beta", [(0.05, 1, 1), (0.1, 5, 2)])
+def test_accept_environment(lam, gamma, beta):
+    DL = (
+        "2 * pi * 2.0 * {lam} / (pi * {gamma} * {beta}) if (w==0) "
+        "else 2 * pi * (2.0 * {lam} * {gamma} * w / (pi * (w**2 + {gamma}**2))) "
+        "* ((1 / (exp(w * {beta}) - 1)) + 1)"
+    ).format(gamma=gamma, beta=beta, lam=lam)
+    H = 0.5 * qutip.sigmax()
+    psi0 = (2 * qutip.basis(2, 0) + qutip.basis(2, 1)).unit()
+    times = np.linspace(0, 10, 100)
+    resultBR_str = brmesolve(
+        H, psi0, times,
+        a_ops=[[qutip.sigmaz(), DL]],
+        e_ops=[qutip.sigmaz()]
+    )
+    env = DrudeLorentzEnvironment(T=1/beta, lam=lam, gamma=gamma)
+    resultBR_env = brmesolve(
+        H, psi0, times,
+        a_ops=[[qutip.sigmaz(), env]],
+        e_ops=[qutip.sigmaz()]
+    )
+    assert np.allclose(resultBR_env.expect[0], resultBR_str.expect[0])