Repository: StingraySoftware/stingray
Pull request #800 (2024-02-12 10:59:15)

Fix deadtime model
The dead time model had a few issues: 

+ the sum of many values below computer precision which produced large errors in the final result -> I set a limit for new additions so that numbers get only added when larger than `10**(-np.finfo(float).precision)`

+ `check_A` and `check_B` had useless axes in some conditions due to the large dynamic range of values -> I changed them to "symlog" to allow for negative numbers with small powers of ten.
![image](https://github.com/StingraySoftware/stingray/assets/7190189/223c6197-2159-4f28-b33f-2ec9bc8f92ce)
I also added a trend line that estimates the estimated floating point error propagation after using the Bks for the actual calculation


Also in this PR: 

+ Added the possibility of calculating `B` and `A` for ranges of the parameter `k`

+ Added a new functions that calculates the deadtime-distorted PDS with more user-friendly parameters (e.g. providing the frequency array, starting from detected (not incident) count rate, using background to rescale, etc.)
![new_deadtime_model](https://github.com/StingraySoftware/stingray/assets/7190189/4d5198df-ef91-49fa-85b9-ddcde227f7ff)

+ Added a method to Crossspectrum and derivatives to directly correct the powers, assuming a constant dead time
![image](https://github.com/StingraySoftware/stingray/assets/7190189/3cd649fa-3bbb-46d4-9e63-8cfa4f8084b9)




codecov[bot]: ## [Codecov](https://app.codecov.io/gh/StingraySoftware/stingray/pull/800?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=StingraySoftware) Report
Attention: `41 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`5285c23`)](https://app.codecov.io/gh/StingraySoftware/stingray/commit/5285c23221ed81c946007e8ae116c84fdc9947d4?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=StingraySoftware) 96.35% compared to head [(`4432094`)](https://app.codecov.io/gh/StingraySoftware/stingray/pull/800?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=StingraySoftware) 78.60%.

| [Files](https://app.codecov.io/gh/StingraySoftware/stingray/pull/800?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=StingraySoftware) | Patch % | Lines |
|---|---|---|
| [stingray/deadtime/model.py](https://app.codecov.io/gh/StingraySoftware/stingray/pull/800?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=StingraySoftware#diff-c3RpbmdyYXkvZGVhZHRpbWUvbW9kZWwucHk=) | 22.64% | [41 Missing :warning: ](https://app.codecov.io/gh/StingraySoftware/stingray/pull/800?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=StingraySoftware) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##             main     #800       +/-   ##
===========================================
- Coverage   96.35%   78.60%   -17.76%     
===========================================
  Files          43       43               
  Lines        8784     8811       +27     
===========================================
- Hits         8464     6926     -1538     
- Misses        320     1885     +1565     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/StingraySoftware/stingray/pull/800?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=StingraySoftware).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=StingraySoftware).

pep8speaks: Hello @matteobachetti! Thanks for updating this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:

* In the file [`stingray/deadtime/model.py`](https://github.com/StingraySoftware/stingray/blob/e308beff395e96b889b61cd46e39751fc032f90e/stingray/deadtime/model.py):

> [Line 19:31](https://github.com/StingraySoftware/stingray/blob/e308beff395e96b889b61cd46e39751fc032f90e/stingray/deadtime/model.py#L19): [E741](https://duckduckgo.com/?q=pep8%20E741) ambiguous variable name 'l'
> [Line 20:21](https://github.com/StingraySoftware/stingray/blob/e308beff395e96b889b61cd46e39751fc032f90e/stingray/deadtime/model.py#L20): [E741](https://duckduckgo.com/?q=pep8%20E741) ambiguous variable name 'l'





Patch:
diff --git a/docs/changes/800.trivial.rst b/docs/changes/800.trivial.rst
new file mode 100644
index 00000000..52fa8810
--- /dev/null
+++ b/docs/changes/800.trivial.rst
@@ -0,0 +1,1 @@
+Dead time model fixes: more stable computations, a friendlier API for the non-paralyzable model, better plotting of check_A and check_B
diff --git a/docs/deadtime.rst b/docs/deadtime.rst
index 13784dd7..85ed2a4c 100644
--- a/docs/deadtime.rst
+++ b/docs/deadtime.rst
@@ -8,5 +8,5 @@ In this tutorial, we will show the effects of dead time on X-ray observations, a
 .. toctree::
    :maxdepth: 2
 
-   notebooks/Deadtime/Check dead time model in Stingray.ipynb
-   notebooks/Deadtime/Check FAD correction in Stingray.ipynb
+   notebooks/Deadtime/Dead time model in Stingray.ipynb
+   notebooks/Deadtime/FAD correction in Stingray.ipynb
diff --git a/docs/notebooks b/docs/notebooks
index f5991ea2..0e34107a 160000
--- a/docs/notebooks
+++ b/docs/notebooks
@@ -1,1 +1,1 @@
-Subproject commit f5991ea2859ca6f32e4e7069a7fba925b5e1f038
+Subproject commit 0e34107ac872db8ed9b4acfcc75b97edf07f3ccb
diff --git a/stingray/crossspectrum.py b/stingray/crossspectrum.py
index 39bc5f03..608dc764 100644
--- a/stingray/crossspectrum.py
+++ b/stingray/crossspectrum.py
@@ -1657,6 +1657,85 @@ class Crossspectrum(StingrayObject):
         self.k = 1
         return
 
+    def deadtime_correct(
+        self, dead_time, rate, background_rate=0, limit_k=200, n_approx=None, paralyzable=False
+    ):
+        """
+        Correct the power spectrum for dead time effects.
+
+        This correction is based on the formula given in Zhang et al. 2015, assuming
+        a constant dead time for all events.
+        For more advanced dead time corrections, see the FAD method from `stingray.deadtime.fad`
+
+        Parameters
+        ----------
+        dead_time: float
+            The dead time of the detector.
+        rate : float
+            Detected source count rate
+
+        Other Parameters
+        ----------------
+        background_rate : float, default 0
+            Detected background count rate. This is important to estimate when deadtime is given by the
+            combination of the source counts and background counts (e.g. in an imaging X-ray detector).
+        paralyzable: bool, default False
+            If True, the dead time correction is done assuming a paralyzable
+            dead time. If False, the correction is done assuming a non-paralyzable
+            (more common) dead time.
+        limit_k : int, default 200
+            Limit to this value the number of terms in the inner loops of
+            calculations. Check the plots returned by  the `check_B` and
+            `check_A` functions to test that this number is adequate.
+        n_approx : int, default None
+            Number of bins to calculate the model power spectrum. If None, it will use the size of
+            the input frequency array. Relatively simple models (e.g., low count rates compared to
+            dead time) can use a smaller number of bins to speed up the calculation, and the final
+            power values will be interpolated.
+
+        Returns
+        -------
+        spectrum: :class:`Crossspectrum` or derivative.
+            The dead-time corrected spectrum.
+        """
+        # I put it here to avoid circular imports
+        from stingray.deadtime import non_paralyzable_dead_time_model
+
+        if paralyzable:
+            raise NotImplementedError("Paralyzable dead time correction is not implemented yet.")
+
+        model = non_paralyzable_dead_time_model(
+            self.freq,
+            dead_time,
+            rate,
+            bin_time=self.dt,
+            limit_k=limit_k,
+            background_rate=background_rate,
+            n_approx=n_approx,
+        )
+        correction = 2 / model
+        new_spec = copy.deepcopy(self)
+        new_spec.power *= correction
+
+        # Now correct internal attributes for both the spectrum and its possible
+        # sub-spectra (PDSs from single channels, dynamical spectra, etc.)
+        objects = [new_spec]
+        if hasattr(new_spec, "pds1"):
+            objects += [new_spec.pds1, new_spec.pds2]
+
+        for obj in objects:
+            for attr in ["unnorm_power", "power_err", "unnorm_power_err"]:
+                if hasattr(obj, attr):
+                    setattr(obj, attr, getattr(obj, attr) * correction)
+
+        for attr in ["cs_all", "unnorm_cs_all"]:
+            if hasattr(new_spec, attr):
+                dynsp = getattr(new_spec, attr)
+                for i, power in enumerate(dynsp):
+                    dynsp[i] = power * correction
+
+        return new_spec
+
 
 class AveragedCrossspectrum(Crossspectrum):
     type = "crossspectrum"
diff --git a/stingray/deadtime/model.py b/stingray/deadtime/model.py
index 8a152947..ce12a3b0 100644
--- a/stingray/deadtime/model.py
+++ b/stingray/deadtime/model.py
@@ -1,143 +1,410 @@
+import warnings
 from stingray.utils import njit, prange
+
 from stingray.loggingconfig import setup_logger
+
+from collections.abc import Iterable
 import numpy as np
 import matplotlib.pyplot as plt
 from scipy.special import factorial
+from scipy.interpolate import interp1d
+
+logger = setup_logger()
 
+MAX_FACTORIAL = 100
+__INVERSE_FACTORIALS = 1.0 / factorial(np.arange(MAX_FACTORIAL))
+PRECISION = np.finfo(float).precision
+TWOPI = np.pi * 2
 
-__FACTORIALS = factorial(np.arange(160))
+STERLING_PARAMETERS = np.array([1 / 12, 1 / 288, -139 / 51840, -571 / 2488320])
 
-logger = setup_logger()
+__all__ = [
+    "r_det",
+    "r_in",
+    "pds_model_zhang",
+    "non_paralyzable_dead_time_model",
+    "check_A",
+    "check_B",
+]
+
+
+@njit()
+def _stirling_factor(m):
+    """First few terms of series expansion appearing in Stirling's approximation."""
+    fact = 1.0
+    power = 1.0
+    for i in range(1, 5):
+        power *= m
+        fact += STERLING_PARAMETERS[i - 1] / power
+    return fact
+
+
+@njit()
+def e_m_x_x_over_factorial(x: float, m: int):
+    r"""Approximate the large number ratio in Eq. 34.
+
+    The original formula for :math:`G_n` (eq. 34 in Zhang+95) has factors of the kind
+    :math:`e^{-x} x^l / l!`. This is the product of very large numbers, that mostly
+    balance each other. In this function, we approximate this product for large :math:`l`
+    starting from Stirling's approximation for the factorial:
+
+    .. math::
+
+       l! \approx \sqrt{2\pi l} (\frac{l}{e})^l A(l)
+
+    where :math:`A(l)` is a series expansion in 1/l of the form
+    :math:`1 + 1 / 12l + 1 / 288l^2 - 139/51840/l^3 + ...`. This allows to transform the product
+    above into
+
+    .. math::
+
+       \frac{e^{-x} x^l}{l!} \approx \frac{1}{A(l)\sqrt{2\pi l}}\left(\frac{x e}{l}\right)^l e^{-x}
+
+    and then, bringing the exponential into the parenthesis
+
+    .. math::
+
+       \frac{e^{-x} x^l}{l!}\approx\frac{1}{A(l)\sqrt{2\pi l}} \left(\frac{x e^{1-x/l}}{l}\right)^l
+
+    The function inside the brackets has a maximum around :math:`x \approx l` and is well-behaved,
+    allowing to approximate the product for large :math:`l` without the need to calculate the
+    factorial or the exponentials directly.
+    """
+    # Numerical errors occur when x is not a float
+    x = float(x)
+
+    if x == 0.0 and m == 0:
+        return 1.0
+
+    # For decently small numbers, use the direct formula
+    if m < 50 or (x > 1 and m * max(np.log10(x), 1) < 50):
+        return np.exp(-x) * x**m * __INVERSE_FACTORIALS[m]
+
+    # Use Stirling's approximation
+    return 1.0 / np.sqrt(TWOPI * m) * np.power(x * np.exp(1 - x / m) / m, m) / _stirling_factor(m)
 
 
 def r_in(td, r_0):
-    """Calculate incident countrate given dead time and detected countrate."""
+    """Calculate incident countrate given dead time and detected countrate.
+
+    Parameters
+    ----------
+    td : float
+        Dead time
+    r_0 : float
+        Detected countrate
+    """
     tau = 1 / r_0
     return 1.0 / (tau - td)
 
 
 def r_det(td, r_i):
-    """Calculate detected countrate given dead time and incident countrate."""
+    """Calculate detected countrate given dead time and incident countrate.
+
+    Parameters
+    ----------
+    td : float
+        Dead time
+    r_i : float
+        Incident countrate
+    """
     tau = 1 / r_i
     return 1.0 / (tau + td)
 
 
 @njit()
-def Gn(x, n):
+def _Gn(x, n):
     """Term in Eq. 34 in Zhang+95."""
-    s = 0
-    for l in range(0, n):
-        s += (n - l) / __FACTORIALS[l] * x**l
-    return np.exp(-x) * s
+    s = 0.0
 
+    for m in range(0, n):
+        new_val = e_m_x_x_over_factorial(x, m) * (n - m)
 
-@njit()
-def heaviside(x):
-    """Heaviside function. Returns 1 if x>0, and 0 otherwise.
-
-    Examples
-    --------
-    >>> heaviside(2)
-    1
-    >>> heaviside(-1)
-    0
-    """
-    if x >= 0:
-        return 1
-    else:
-        return 0
+        s += new_val
+        # The curve above has a maximum around x~l
+        if x != 0 and s > 0 and m > 2 * x and -np.log10(np.abs(new_val / s)) > PRECISION:
+            break
+
+    return s
 
 
 @njit()
-def h(k, n, td, tb, tau):
+def _h(k, n, td, tb, tau):
     """Term in Eq. 35 in Zhang+95."""
     # Typo in Zhang+95 corrected. k * tb, not k * td
-    if k * tb < n * td:
-        return 0
-    return k - n * (td + tau) / tb + tau / tb * Gn((k * tb - n * td) / tau, n)
+    factor = k * tb - n * td
+
+    if k * tb - n * td < 0:
+        return 0.0
+
+    val = k - n * (td + tau) / tb + tau / tb * _Gn(factor / tau, n)
+
+    return val
 
 
-INFINITE = 100
+INFINITE = 699
 
 
 @njit()
 def A0(r0, td, tb, tau):
-    """Term in Eq. 38 in Zhang+95."""
-    s = 0
-    for n in range(1, INFINITE):
-        s += h(1, n, td, tb, tau)
+    """Term in Eq. 38 in Zhang+95.
+
+    Parameters
+    ----------
+    r0 : float
+        Detected countrate
+    td : float
+        Dead time
+    tb : float
+        Bin time of the light curve
+    tau : float
+        Inverse of the incident countrate
+    """
+    s = 0.0
+    for n in range(1, int(max(2, tb / td * 2 + 1))):
+        s += _h(1, n, td, tb, tau)
 
     return r0 * tb * (1 + 2 * s)
 
 
 @njit()
-def A(k, r0, td, tb, tau):
-    """Term in Eq. 39 in Zhang+95."""
+def A_single_k(k, r0, td, tb, tau):
+    """Term in Eq. 39 in Zhang+95.
+
+    Parameters
+    ----------
+    k : int
+        Order of the term
+    r0 : float
+        Detected countrate
+    td : float
+        Dead time
+    tb : float
+        Bin time of the light curve
+    tau : float
+        Inverse of the incident countrate
+    """
     if k == 0:
         return A0(r0, td, tb, tau)
     # Equation 39
-    s = 0
-    for n in range(1, INFINITE):
-        s += h(k + 1, n, td, tb, tau) - 2 * h(k, n, td, tb, tau) + h(k - 1, n, td, tb, tau)
+    s = 0.0
+
+    for n in range(1, int(max(3, (k + 1) * tb / td * 2))):
+        new_val = _h(k + 1, n, td, tb, tau) - 2 * _h(k, n, td, tb, tau) + _h(k - 1, n, td, tb, tau)
+
+        s += new_val
 
     return r0 * tb * s
 
 
-def check_A(rate, td, tb, max_k=100, save_to=None):
-    """Test that A is well-behaved.
+def A(k, r0, td, tb, tau):
+    """Term in Eq. 39 in Zhang+95.
+
+    Parameters
+    ----------
+    k : int or array of ints
+        Order of the term
+    r0 : float
+        Detected countrate
+    td : float
+        Dead time
+    tb : float
+        Bin time of the light curve
+    tau : float
+        Inverse of the incident countrate
+    """
+    if isinstance(k, Iterable):
+        return np.array([A_single_k(ki, r0, td, tb, tau) for ki in k])
+
+    return A_single_k(k, r0, td, tb, tau)
+
+
+def limit_A(rate, td, tb):
+    """Limit of A for k->infty, as per Eq. 43 in Zhang+95.
 
-    Check that Ak ->r0**2tb**2 for k->infty, as per Eq. 43 in
-    Zhang+95.
+    Parameters
+    ----------
+    rate : float
+        Incident count rate
+    td : float
+        Dead time
+    tb : float
+        Bin time of the light curve
     """
-    tau = 1 / rate
     r0 = r_det(td, rate)
+    return r0**2 * tb**2
 
-    value = r0**2 * tb**2
+
+def check_A(rate, td, tb, max_k=100, save_to=None, linthresh=0.000001, rate_is_incident=True):
+    r"""Test that A is well-behaved.
+
+    This function produces a plot of :math:`A_k - r_0^2 t_b^2` vs :math:`k`, to visually check that
+    :math:`A_k \rightarrow r_0^2 t_b^2` for :math:`k\rightarrow\infty`, as per Eq. 43 in Zhang+95.
+
+    With this function is possible to determine how many inner loops `k` (`limit_k` in function
+    pds_model_zhang) are necessary for a correct approximation of the dead time model
+
+    Parameters
+    ----------
+    rate : float
+        Count rate, either incident or detected (use the `rate_is_incident` bool to specify)
+    td : float
+        Dead time
+    tb : float
+        Bin time of the light curve
+
+    Other Parameters
+    ----------------
+    max_k : int
+        Maximum k to plot
+    save_to : str, default None
+        If not None, save the plot to this file
+    linthresh : float, default 0.000001
+        Linear threshold for the "symlog" scale of the plot
+    rate_is_incident : bool, default True
+        If True, the input rate is the incident count rate. If False, it is the detected one.
+    """
+    if rate_is_incident:
+        tau = 1 / rate
+        r0 = r_det(td, rate)
+    else:
+        r0 = rate
+        tau = 1 / r_in(td, rate)
+
+    limit = limit_A(rate, td, tb)
     fig = plt.figure()
-    for k in range(max_k):
-        plt.scatter(k, A(k, r0, td, tb, tau), color="k")
-    plt.axhline(value, ls="--", color="k")
+
+    k_values = np.arange(0, max_k + 1)
+    A_values = A(k_values, r0, td, tb, tau)
+
+    plt.plot(k_values, A_values - limit, color="k")
+    plt.semilogx()
+    plt.yscale("symlog", linthresh=linthresh)
+    plt.axhline(0, ls="--", color="k")
     plt.xlabel("$k$")
-    plt.ylabel("$A_k$")
+    plt.ylabel("$A_k - r_0^2 t_b^2$")
     if save_to is not None:
         plt.savefig(save_to)
         plt.close(fig)
+    return k_values, A_values
 
 
 @njit()
-def B(k, r0, td, tb, tau):
+def _B_raw(k, r0, td, tb, tau):
     """Term in Eq. 45 in Zhang+95."""
     if k == 0:
-        return 2 * (A(0, r0, td, tb, tau) - r0**2 * tb**2) / (r0 * tb)
+        return 2 * (A_single_k(0, r0, td, tb, tau) - r0**2 * tb**2) / (r0 * tb)
+
+    new_val = A_single_k(k, r0, td, tb, tau) - r0**2 * tb**2
 
-    return 4 * (A(k, r0, td, tb, tau) - r0**2 * tb**2) / (r0 * tb)
+    return 4 * new_val / (r0 * tb)
 
 
 @njit()
-def safe_B(k, r0, td, tb, tau, limit_k=60):
+def _safe_B_single_k(k, r0, td, tb, tau, limit_k=60):
     """Term in Eq. 39 in Zhang+95, with a cut in the maximum k.
 
     This can be risky. Only use if B is really 0 for high k.
     """
     if k > limit_k:
-        return 0
-    return B(k, r0, td, tb, tau)
+        return 0.0
+    return _B_raw(k, r0, td, tb, tau)
 
 
-def check_B(rate, td, tb, max_k=100, save_to=None):
-    """Check that B->0 for k->infty."""
-    tau = 1 / rate
-    r0 = r_det(td, rate)
+def _safe_B(k, r0, td, tb, tau, limit_k=60):
+    """Term in Eq. 39 in Zhang+95, with a cut in the maximum k.
+
+    This can be risky. Only use if B is really 0 for high k.
+    """
+    if isinstance(k, Iterable):
+        return np.array([_safe_B_single_k(ki, r0, td, tb, tau, limit_k=limit_k) for ki in k])
+    return _safe_B_single_k(int(k), r0, td, tb, tau, limit_k=limit_k)
+
+
+def B(k, r0, td, tb, tau, limit_k=60):
+    """Term in Eq. 39 in Zhang+95, with a cut in the maximum k.
+
+    The cut can be risky. Only use if B is really 0 for high k. Use `check_B` to test.
+
+    Parameters
+    ----------
+    k : int or array of ints
+        Order of the term
+    r0 : float
+        Detected countrate
+    td : float
+        Dead time
+    tb : float
+        Bin time of the light curve
+    tau : float
+        Inverse of the incident countrate
+
+    Other Parameters
+    ----------------
+    limit_k : int
+        Limit to this value the number of terms in the inner loops of
+        calculations. Check the plots returned by  the `check_B` and
+        `check_A` functions to test that this number is adequate.
+    """
+    return _safe_B(k, r0, td, tb, tau, limit_k=limit_k)
+
+
+def check_B(rate, td, tb, max_k=100, save_to=None, linthresh=0.000001, rate_is_incident=True):
+    r"""Check that :math:`B\rightarrow 0` for :math:`k\rightarrow \infty`.
+
+    This function produces a plot of :math:`B_k` vs :math:`k`, to visually check that
+    :math:`B_k \rightarrow 0` for :math:`k\rightarrow\infty`, as per Eq. 43 in Zhang+95.
+
+    With this function is possible to determine how many inner loops `k` (`limit_k` in function
+    pds_model_zhang) are necessary for a correct approximation of the dead time model
+
+    Parameters
+    ----------
+    rate : float
+        Count rate, either incident or detected (use the `rate_is_incident` bool to specify)
+    td : float
+        Dead time
+    tb : float
+        Bin time of the light curve
+
+    Other Parameters
+    ----------------
+    max_k : int
+        Maximum k to plot
+    save_to : str, default None
+        If not None, save the plot to this file
+    linthresh : float, default 0.000001
+        Linear threshold for the "symlog" scale of the plot
+    rate_is_incident : bool, default True
+        If True, the input rate is the incident count rate. If False, it is the detected one.
+    """
+    if rate_is_incident:
+        tau = 1 / rate
+        r0 = r_det(td, rate)
+    else:
+        r0 = rate
+        tau = 1 / r_in(td, rate)
 
     fig = plt.figure()
-    for k in range(max_k):
-        plt.scatter(k, B(k, r0, td, tb, tau), color="k")
+    k_values = np.arange(1, max_k + 1)
+    B_values = B(k_values, r0, td, tb, tau, limit_k=max_k)
+    threshold = max(tb / td, 1) * 10.0 ** (3 * np.log10(k_values) - PRECISION)
+
+    plt.plot(k_values, threshold, label="Est. propagation of float errors", color="red")
+    plt.plot(k_values, -threshold, color="r")
+
+    plt.plot(k_values, B_values, color="k", ds="steps-mid")
     plt.axhline(0, ls="--", color="k")
+    plt.semilogx()
+    plt.yscale("symlog", linthresh=linthresh)
     plt.xlabel("$k$")
     plt.ylabel("$B_k$")
+    plt.xlim([0.9, max_k])
+    plt.legend()
     if save_to is not None:
         plt.savefig(save_to)
         plt.close(fig)
+    return k_values, B_values
 
 
 @njit(parallel=True)
@@ -145,21 +412,19 @@ def _inner_loop_pds_zhang(N, tau, r0, td, tb, limit_k=60):
     """Calculate the power spectrum, as per Eq. 44 in Zhang+95."""
     P = np.zeros(N // 2)
     for j in prange(N // 2):
-        eq8_sum = 0
-        for k in range(1, N):
-            eq8_sum += (
-                (N - k)
-                / N
-                * safe_B(k, r0, td, tb, tau, limit_k=limit_k)
-                * np.cos(2 * np.pi * j * k / N)
-            )
+        eq8_sum = 0.0
+
+        for k in range(1, min(N, limit_k)):
+            Bk = _safe_B_single_k(k, r0, td, tb, tau, limit_k=limit_k)
+
+            eq8_sum += (N - k) / N * Bk * np.cos(2 * np.pi * j * k / N)
 
-        P[j] = safe_B(0, r0, td, tb, tau) + eq8_sum
+        P[j] = _safe_B_single_k(0, r0, td, tb, tau) + eq8_sum
 
     return P
 
 
-def pds_model_zhang(N, rate, td, tb, limit_k=60):
+def pds_model_zhang(N, rate, td, tb, limit_k=60, rate_is_incident=True):
     """Calculate the dead-time-modified power spectrum.
 
     Parameters
@@ -179,6 +444,9 @@ def pds_model_zhang(N, rate, td, tb, limit_k=60):
         Limit to this value the number of terms in the inner loops of
         calculations. Check the plots returned by  the `check_B` and
         `check_A` functions to test that this number is adequate.
+    rate_is_incident : bool, default True
+        If True, the input rate is the incident count rate. If False, it is the
+        detected count rate.
 
     Returns
     -------
@@ -187,14 +455,94 @@ def pds_model_zhang(N, rate, td, tb, limit_k=60):
     power : array of floats
         Power spectrum
     """
-    tau = 1 / rate
-    r0 = r_det(td, rate)
+    if rate_is_incident:
+        tau = 1 / rate
+        r0 = r_det(td, rate)
+    else:
+        r0 = rate
+        tau = 1 / r_in(td, rate)
+
     # Nph = N / tau
     logger.info("Calculating PDS model (update)")
     P = _inner_loop_pds_zhang(N, tau, r0, td, tb, limit_k=limit_k)
+    if tb > 10 * td:
+        warnings.warn(
+            f"The bin time is much larger than the dead time. "
+            f" Calculations might be slow. tb={tb / td:.2f} * td"
+        )
 
     maxf = 0.5 / tb
     df = maxf / len(P)
-    freqs = np.arange(0, maxf, df)
+    freqs = np.arange(len(P)) * df
 
     return freqs, P
+
+
+def non_paralyzable_dead_time_model(
+    freqs,
+    dead_time,
+    rate,
+    bin_time=None,
+    limit_k=200,
+    background_rate=0.0,
+    n_approx=None,
+):
+    """Calculate the dead-time-modified power spectrum.
+
+    Parameters
+    ----------
+    freqs : array of floats
+        Frequency array
+    dead_time : float
+        Dead time
+    rate : float
+        Detected source count rate
+
+    Other Parameters
+    ----------------
+    bin_time : float
+        Bin time of the light curve
+    limit_k : int, default 200
+        Limit to this value the number of terms in the inner loops of
+        calculations. Check the plots returned by  the `check_B` and
+        `check_A` functions to test that this number is adequate.
+    background_rate : float, default 0
+        Detected background count rate. This is important to estimate when deadtime is given by the
+        combination of the source counts and background counts (e.g. in an imaging X-ray detector).
+    n_approx : int, default None
+        Number of bins to calculate the model power spectrum. If None, it will use the size of
+        the input frequency array. Relatively simple models (e.g., low count rates compared to
+        dead time) can use a smaller number of bins to speed up the calculation, and the final
+        power values will be interpolated.
+
+    Returns
+    -------
+    power : array of floats
+        Power spectrum
+    """
+    if rate + background_rate > 1 / dead_time:
+        raise ValueError(
+            "The sum of the source and background count rates is larger than the inverse "
+            "of the dead time. This is not a physical situation. Please check your input."
+        )
+
+    if bin_time is None:
+        bin_time = 1 / (2 * max(freqs))
+
+    n_bins = n_approx if n_approx is not None else max(max(freqs), 10)
+
+    zh_f, zh_p = pds_model_zhang(
+        int(n_bins) * 2,
+        (rate + background_rate),
+        dead_time,
+        bin_time,
+        limit_k=limit_k,
+        rate_is_incident=False,
+    )
+
+    # Rescale by the source rate wrt background rate
+    if background_rate > 0:
+        zh_p = (zh_p - 2) * rate / (rate + background_rate) + 2
+
+    pds_interp = interp1d(zh_f, zh_p, bounds_error=False, fill_value="extrapolate")
+    return pds_interp(freqs)
diff --git a/stingray/utils.py b/stingray/utils.py
index 80090589..264fc17d 100644
--- a/stingray/utils.py
+++ b/stingray/utils.py
@@ -140,9 +140,27 @@ __all__ = [
     "nearest_power_of_two",
     "find_nearest",
     "check_isallfinite",
+    "heaviside",
 ]
 
 
+@njit()
+def heaviside(x):
+    """Heaviside function. Returns 1 if x>0, and 0 otherwise.
+
+    Examples
+    --------
+    >>> heaviside(2)
+    1
+    >>> heaviside(-1)
+    0
+    """
+    if x >= 0:
+        return 1
+    else:
+        return 0
+
+
 @njit
 def any_complex_in_array(array):
     """Check if any element of an array is complex.


Test patch:
diff --git a/stingray/deadtime/tests/test_models.py b/stingray/deadtime/tests/test_models.py
index 89b94d85..4dce60f0 100644
--- a/stingray/deadtime/tests/test_models.py
+++ b/stingray/deadtime/tests/test_models.py
@@ -5,24 +5,25 @@ from scipy.interpolate import interp1d
 
 from stingray.lightcurve import Lightcurve
 from stingray.powerspectrum import AveragedPowerspectrum
-from stingray.deadtime.model import r_det, r_in, pds_model_zhang
-from stingray.deadtime.model import check_A, check_B, heaviside
+from stingray.deadtime.model import r_det, r_in, pds_model_zhang, non_paralyzable_dead_time_model
+from stingray.deadtime.model import check_A, check_B, A, B
 from stingray.filters import filter_for_deadtime
+from stingray.utils import HAS_NUMBA
 
 
 pytestmark = pytest.mark.slow
 
 
-def test_heaviside():
-    assert heaviside(2) == 1
-    assert heaviside(0) == 1
-    assert heaviside(-1) == 0
-
-
-def simulate_events(rate, length, deadtime=2.5e-3, **filter_kwargs):
+def simulate_events(rate, length, deadtime=2.5e-3, bkg_rate=0.0, **filter_kwargs):
     events = np.random.uniform(0, length, int(rate * length))
+
     events = np.sort(events)
-    events_dt = filter_for_deadtime(events, deadtime, **filter_kwargs)
+
+    events_back = None
+    if bkg_rate > 0:
+        events_back = np.sort(np.random.uniform(0, length, int(rate * length)))
+
+    events_dt = filter_for_deadtime(events, deadtime, bkg_ev_list=events_back, **filter_kwargs)
     return events, events_dt
 
 
@@ -35,32 +36,85 @@ def test_deadtime_conversion():
     np.testing.assert_almost_equal(rin, original_rate)
 
 
-def test_zhang_model_accurate():
-    bintime = 1 / 4096
+@pytest.mark.parametrize("rate", [1.0, 100.0])
+def test_zhang_model_accurate(rate):
+    bintime = 0.0002
     deadtime = 2.5e-3
     length = 2000
-    fftlen = 5
-    r = 300
+    fftlen = 10
 
-    events, events_dt = simulate_events(r, length, deadtime=deadtime)
+    _, events_dt = simulate_events(rate, length, deadtime=deadtime)
     lc_dt = Lightcurve.make_lightcurve(events_dt, bintime, tstart=0, tseg=length)
     pds = AveragedPowerspectrum(lc_dt, fftlen, norm="leahy")
 
-    zh_f, zh_p = pds_model_zhang(1000, r, deadtime, bintime, limit_k=100)
+    zh_f, zh_p = pds_model_zhang(100, rate, deadtime, bintime, limit_k=600)
 
     deadtime_fun = interp1d(zh_f, zh_p, bounds_error=False, fill_value="extrapolate")
     ratio = pds.power / deadtime_fun(pds.freq)
-    assert np.isclose(np.mean(ratio), 1, atol=0.001)
-    assert np.isclose(np.std(ratio), 1 / np.sqrt(pds.m), atol=0.001)
+    assert np.isclose(np.mean(ratio), 1, rtol=0.01)
+    assert np.isclose(np.std(ratio), 1 / np.sqrt(pds.m), rtol=0.1)
+
 
+@pytest.mark.skipif("not HAS_NUMBA")
+@pytest.mark.parametrize("rates", [(1.0, 0.0), (1.0, 1.0), (100.0, 10.0), (100, 200)])
+def test_non_paralyzable_model_accurate(rates):
+    bintime = 0.0002
+    deadtime = 2.5e-3
+    length = 2000
+    fftlen = 10
+    rate, bkg_rate = rates
+    events, events_dt = simulate_events(rate, length, deadtime=deadtime, bkg_rate=bkg_rate)
+    det_rate = events_dt.size / length
+    det_bkg_rate = events.size / length - det_rate
 
-def test_checkA():
-    check_A(300, 2.5e-3, 0.001, max_k=100, save_to="check_A.png")
+    lc_dt = Lightcurve.make_lightcurve(events_dt, bintime, tstart=0, tseg=length)
+    pds = AveragedPowerspectrum(lc_dt, fftlen, norm="leahy")
+
+    model_power = non_paralyzable_dead_time_model(
+        pds.freq, deadtime, rate=det_rate, background_rate=det_bkg_rate
+    )
+    ratio = pds.power / model_power
+    assert np.isclose(np.mean(ratio), 1, rtol=0.01)
+
+
+@pytest.mark.parametrize("is_incident", [True, False])
+def test_checkA(is_incident):
+    check_A(300, 2.5e-3, 0.001, max_k=100, save_to="check_A.png", rate_is_incident=is_incident)
     assert os.path.exists("check_A.png")
     os.unlink("check_A.png")
 
 
-def test_checkB():
-    check_B(300, 2.5e-3, 0.001, max_k=100, save_to="check_B.png")
+@pytest.mark.parametrize("is_incident", [True, False])
+def test_checkB(is_incident):
+    check_B(300, 2.5e-3, 0.001, max_k=100, save_to="check_B.png", rate_is_incident=is_incident)
     assert os.path.exists("check_B.png")
     os.unlink("check_B.png")
+
+
+@pytest.mark.parametrize("tb", [0.0001, 0.1])
+def test_A_and_B_array(tb):
+    td = 2.5e-3
+    ks = np.array([1, 5, 20, 70])
+    rate = 10
+    tau = 1 / rate
+    r0 = r_det(td, rate)
+    assert np.array_equal(
+        np.array([A(k, r0, td, tb, tau) for k in ks]), A(ks, r0, td, tb, tau), equal_nan=False
+    )
+    assert np.array_equal(
+        np.array([B(k, r0, td, tb, tau) for k in ks]), B(ks, r0, td, tb, tau), equal_nan=False
+    )
+
+
+def test_pds_model_warns():
+    with pytest.warns(UserWarning, match="The bin time is much larger than the "):
+        pds_model_zhang(10, 100.0, 2.5e-3, 0.1, limit_k=10)
+
+
+def test_non_paralyzable_model_fail_bad_rate():
+    """Fail if combined rate is larger than 1 / deadtime."""
+    with pytest.raises(
+        ValueError,
+        match="The sum of the source and background count rates is larger than the inverse",
+    ):
+        non_paralyzable_dead_time_model(np.arange(10), 2.5e-3, rate=300, background_rate=300)
diff --git a/stingray/tests/test_crossspectrum.py b/stingray/tests/test_crossspectrum.py
index 90f6741e..b7e1919a 100644
--- a/stingray/tests/test_crossspectrum.py
+++ b/stingray/tests/test_crossspectrum.py
@@ -9,12 +9,14 @@ from astropy.io import fits
 from stingray import Lightcurve
 from stingray import Crossspectrum, AveragedCrossspectrum, DynamicalCrossspectrum
 from stingray import AveragedPowerspectrum
-from stingray.crossspectrum import cospectra_pvalue
+from stingray.crossspectrum import cospectra_pvalue, crossspectrum_from_time_array
 from stingray.crossspectrum import normalize_crossspectrum, normalize_crossspectrum_gauss
 from stingray.crossspectrum import coherence, time_lag
 from stingray import StingrayError
+from stingray.utils import HAS_NUMBA
 from stingray.simulator import Simulator
 from stingray.fourier import poisson_level
+from stingray.filters import filter_for_deadtime
 
 from stingray.events import EventList
 import copy
@@ -1249,6 +1251,33 @@ class TestAveragedCrossspectrum(object):
         maxpower = np.max(cs.power)
         assert np.all(np.isfinite(cs.classical_significances(threshold=maxpower / 2.0)))
 
+    @pytest.mark.skipif("not HAS_NUMBA")
+    def test_deadtime_corr(self):
+        tmax = 100.0
+        segment_size = 1
+        events1 = np.sort(np.random.uniform(0, tmax, 10000))
+        events2 = np.sort(np.random.uniform(0, tmax, 10000))
+        events1_dt = filter_for_deadtime(events1, deadtime=0.0025)
+        events2_dt = filter_for_deadtime(events2, deadtime=0.0025)
+        cs_dt = crossspectrum_from_time_array(
+            events1_dt,
+            events2_dt,
+            gti=[[0, tmax]],
+            dt=0.001,
+            segment_size=segment_size,
+            save_all=True,
+            norm="leahy",
+        )
+        # Paralyzable is not implemented yet
+        with pytest.raises(NotImplementedError):
+            cs = cs_dt.deadtime_correct(
+                dead_time=0.0025, rate=np.size(events1_dt) / tmax, paralyzable=True
+            )
+
+        cs = cs_dt.deadtime_correct(dead_time=0.0025, rate=np.size(events1_dt) / tmax)
+        # Poisson noise has a scatter of sqrt(2/N) in the cospectrum
+        assert np.isclose(np.std(cs.power.real), np.sqrt(2 / (tmax / segment_size)), rtol=0.1)
+
 
 class TestCoherenceFunction(object):
     def setup_class(self):
diff --git a/stingray/tests/test_powerspectrum.py b/stingray/tests/test_powerspectrum.py
index 13ddca43..20549326 100644
--- a/stingray/tests/test_powerspectrum.py
+++ b/stingray/tests/test_powerspectrum.py
@@ -9,8 +9,11 @@ import matplotlib.pyplot as plt
 from astropy.io import fits
 from stingray import Lightcurve
 from stingray.events import EventList
+from stingray.utils import HAS_NUMBA
 from stingray import Powerspectrum, AveragedPowerspectrum, DynamicalPowerspectrum
+from stingray.powerspectrum import powerspectrum_from_time_array
 from astropy.modeling.models import Lorentz1D
+from stingray.filters import filter_for_deadtime
 
 _HAS_XARRAY = importlib.util.find_spec("xarray") is not None
 _HAS_PANDAS = importlib.util.find_spec("pandas") is not None
@@ -340,6 +343,25 @@ class TestAveragedPowerspectrumEvents(object):
         assert np.isclose(np.mean(ps.power), 2.0, atol=1e-2, rtol=1e-2)
         assert np.isclose(np.std(ps.power), 2.0 / np.sqrt(ps.m), atol=0.1, rtol=0.1)
 
+    @pytest.mark.skipif("not HAS_NUMBA")
+    def test_deadtime_corr(self):
+        tmax = 100.0
+        segment_size = 1
+        events = np.sort(np.random.uniform(0, tmax, 10000))
+        events_dt = filter_for_deadtime(events, deadtime=0.0025)
+        pds_dt = powerspectrum_from_time_array(
+            events_dt,
+            gti=[[0, tmax]],
+            dt=0.001,
+            segment_size=segment_size,
+            save_all=True,
+            norm="leahy",
+        )
+
+        pds = pds_dt.deadtime_correct(dead_time=0.0025, rate=np.size(events_dt) / tmax)
+        assert np.isclose(np.mean(pds.power), 2, rtol=0.1)
+        assert np.isclose(np.std(pds.power), 2 / np.sqrt(tmax / segment_size), rtol=0.1)
+
 
 class TestPowerspectrum(object):
     @classmethod
@@ -481,9 +503,7 @@ class TestPowerspectrum(object):
         square of the number of data points in the light curve
         """
         ps = Powerspectrum(self.lc, norm="Leahy")
-        ps_var = (np.sum(self.lc.counts) / ps.n**2.0) * (
-            np.sum(ps.power[:-1]) + ps.power[-1] / 2.0
-        )
+        ps_var = (np.sum(self.lc.counts) / ps.n**2) * (np.sum(ps.power[:-1]) + ps.power[-1] / 2.0)
 
         assert np.isclose(ps_var, np.var(self.lc.counts), atol=0.01)
 
diff --git a/stingray/tests/test_utils.py b/stingray/tests/test_utils.py
index 3f674967..bc4589ad 100644
--- a/stingray/tests/test_utils.py
+++ b/stingray/tests/test_utils.py
@@ -367,6 +367,12 @@ def test_equal_count_energy_ranges():
         assert np.count_nonzero(good) == 10000
 
 
+def test_heaviside():
+    assert utils.heaviside(2) == 1
+    assert utils.heaviside(0) == 1
+    assert utils.heaviside(-1) == 0
+
+
 def test_histogram_equiv_numpy():
     x = np.random.uniform(0.0, 1.0, 100)
     (H, _) = np.histogram(x, bins=5, range=(0.0, 1.0))