made tests faster

qutip/core/environment.py

@@ -1059,9 +1059,7 @@ class UnderDampedEnvironment(BosonicEnvironment):
     r"""
     Describes an underdamped environment with the spectral density
 
-    .. math::
-        J(\omega) = \frac{\lambda^{2} \Gamma \omega}{(\omega_{c}^{2}-
-        \omega^{2})^{2}+ \Gamma^{2} \omega^{2}}
+
 
     (see Eq. 16 in [BoFiN23]_).
 

qutip/tests/core/test_environment.py

@@ -8,6 +8,55 @@
 )
 from scipy.integrate import quad_vec
 
+# Tests are slow if precise, so tolerance is slow for speed
+# tolerance should be one order of magnitude below the epsabs,epsrel in
+# correlation
+inttol = 1e-3
+comtol = 1e-2
+
+
+def spectral_density(w, lam, wc):
+    return lam*w*np.exp(-np.abs(w)/wc)
+
+
+def power(w, lam, wc, T):
+    zero_mask = (w == 0)
+    nonzero_mask = np.invert(zero_mask)
+
+    S = np.zeros_like(w)
+    S[zero_mask] = 2 * T * spectral_density(1e-16, lam, wc) / 1e-16
+    S[nonzero_mask] = 2*np.sign(w[nonzero_mask])*spectral_density(
+        np.abs(w[nonzero_mask]), lam, wc)*(
+        n_thermal(w[nonzero_mask], T)+1)
+    return S
+
+
+def correlation(t, lam, wc, T):
+    def integrand(w, t):
+        return spectral_density(w, lam, wc) / np.pi * (
+            (2 * n_thermal(w, T) + 1) * np.cos(w * t)
+            - 1j * np.sin(w * t)
+        )
+
+    result = quad_vec(lambda w: integrand(w, t), 0,
+                      np.Inf, epsabs=inttol, epsrel=inttol)
+    return result[0]
+
+
+@pytest.fixture()
+def params():
+    N = 3
+    np.random.seed(42)
+    lam = np.random.uniform(low=0.05, high=0.5, size=(N,))
+    wc = np.random.uniform(low=2, high=5, size=(N,))
+    T = np.random.uniform(low=1, high=5, size=(N,))
+    t = np.linspace(0, 25, 100)
+    w = np.linspace(0, 15, 100)
+    w2 = np.linspace(-20, 20, 100)
+    corr = [correlation(t, wc[k], lam[k], T[k])
+            for k in range(len(lam))]
+    return lam, wc, T, t, w, w2, corr
+
 
 def spectral_density(w, gamma, w0, lam):
     return lam**2 * gamma * w / ((w**2 - w0**2)**2
@@ -41,13 +90,14 @@ def integrand(w, t):
 @pytest.fixture()
 def params():
     N = 3
-    lam = np.random.uniform(low=0.05, high=0.5, size=(N,))
-    gamma = np.random.uniform(low=2, high=5, size=(N,))
+    np.random.seed(42)
+    lam = np.random.uniform(low=0.1, high=0.9, size=(N,))
+    gamma = np.random.uniform(low=1, high=2, size=(N,))
     w0 = np.random.uniform(low=1.1, high=5, size=(N,))
-    T = np.random.uniform(low=1, high=5, size=(N,))
-    t = np.linspace(0, 25, 500)
-    w = np.linspace(0, 25, 1000)
-    w2 = np.linspace(-50, 50, 1000)
+    T = np.random.uniform(low=1.5, high=5, size=(N,))
+    t = np.linspace(0, 10, 500)
+    w = np.linspace(0, 25, 500)
+    w2 = np.linspace(-50, 50, 500)
     corr = [correlation(t, gamma[k], w0[k], lam[k], T[k])
             for k in range(len(lam))]
     return lam, gamma, w0, T, t, w, w2, corr
@@ -70,7 +120,7 @@ def test_from_correlation(self, params):
                 w, gamma[k], w0[k], lam[k]), atol=1e-2).all()
 
     def test_from_spectral_density(self, params):
-        lam, gamma, w0, T, t, w, w2, corr = params
+        lam, gamma, w0, T, t, w, w2, _ = params
         for k in range(len(lam)):
             sd = spectral_density(w, gamma[k], w0[k], lam[k])
             bb5 = BosonicEnvironment.from_spectral_density(
@@ -87,46 +137,49 @@ def test_from_power_spectrum(self, params):
         for k in range(len(lam)):
             pow = power(w2, gamma[k], w0[k], lam[k], T[k])
             bb5 = BosonicEnvironment.from_power_spectrum(
-                pow, w2, wMax=10*gamma[k], T=T[k])
-            assert np.isclose(bb5.correlation_function(t), correlation(
-                t, gamma[k], w0[k], lam[k], T[k]), atol=1e-2).all()
+                pow, w2, wMax=gamma[k], T=T[k])
+            assert np.isclose(bb5.correlation_function(t),
+                              corr[k], atol=comtol).all()
             assert np.isclose(bb5.power_spectrum(w2), power(
                 w2, gamma[k], w0[k], lam[k], T[k])).all()
             assert np.isclose(bb5.spectral_density(w), spectral_density(
-                w, gamma[k], w0[k], lam[k]), atol=1e-2).all()
+                w, gamma[k], w0[k], lam[k]), atol=comtol).all()
 
     def test_approx_by_cf_fit(self, params):
         lam, gamma, w0, T, t, w, w2, corr = params
         for k in range(len(lam)):
             bb5 = BosonicEnvironment.from_correlation_function(
                 corr[k], t, T=T[k])
-            bb6, _ = bb5.approx_by_cf_fit(t, target_rsme=2e-5)
+            bb6, finfo = bb5.approx_by_cf_fit(
+                t, target_rsme=None, Nr_max=2, Ni_max=1)
+            print(finfo['summary'])
             assert np.isclose(bb6.correlation_function(t),
-                              corr[k], atol=1e-3).all()
+                              corr[k], atol=5*comtol).all()
             assert np.isclose(bb6.power_spectrum(w2), power(
-                w2, gamma[k], w0[k], lam[k], T[k]), atol=1e-2).all()
+                w2, gamma[k], w0[k], lam[k], T[k]), atol=5*comtol).all()
             assert np.isclose(bb6.spectral_density(w), spectral_density(
-                w, gamma[k], w0[k], lam[k]), atol=1e-2).all()
+                w, gamma[k], w0[k], lam[k]), atol=5*comtol).all()
 
     def test_approx_by_sd_fit(self, params):
         lam, gamma, w0, T, t, w, w2, corr = params
         for k in range(len(lam)):
             sd = spectral_density(t, gamma[k], w0[k], lam[k])
             bb5 = BosonicEnvironment.from_spectral_density(sd, t, T=T[k])
-            bb6, finfo = bb5.approx_by_sd_fit(w, target_rsme=1e-4)
+            bb6, finfo = bb5.approx_by_sd_fit(w, Nmax=1, Nk=10)
+            print(gamma[k], lam[k], w0[k])
+            print(finfo['summary'])
             # asking for more precision that the Bosonic enviroment has may
             # violate this (due to the interpolation and it's easily fixed
             # using a denser range). I could have set N=1 but thought this was
             # a sensible test for N
-            print(gamma[k], w0[k], lam[k])
             assert finfo["N"] == 1
             assert np.isclose(bb6.correlation_function(t),
-                              correlation(t, gamma[k], w0[k], lam[k], T[k]),
-                              atol=1e-2).all()
+                              corr[k],
+                              atol=comtol).all()
             assert np.isclose(bb6.power_spectrum(w2), power(
-                w2, gamma[k], w0[k], lam[k], T[k]), atol=1e-2).all()
+                w2, gamma[k], w0[k], lam[k], T[k]), atol=comtol).all()
             assert np.isclose(bb6.spectral_density(w), spectral_density(
-                w, gamma[k], w0[k], lam[k]), atol=1e-2).all()
+                w, gamma[k], w0[k], lam[k]), atol=comtol).all()
 
 
 class TestFits: