Compute relative error for the relativistic Breit-Wigner PDF calculation.

check_rel_breitwigner_pdf.py

from mpmath import mp
from mpsci.distributions import rel_breitwigner

import numpy as np
from scipy import stats
from scipy.special import powm1

import matplotlib.pyplot as plt


def _pdf_v2(x, rho):
    # This is a slight variation of _pdf() from the PR.
    # C = k / rho**2
    C = np.sqrt(
        2 * (1 + 1/rho**2) / (1 + np.sqrt(1 + 1/rho**2))
    ) * 2 / np.pi
    # return C / (((x**2 - rho**2)/rho)**2 + 1)
    with np.errstate(over='ignore'):
        return C / (((x - rho)*(x + rho)/rho)**2 + 1)


def _k(rho):
    rho2 = rho**2
    s = np.sqrt(rho2 + 1)
    return 2*np.sqrt(2)*rho2*s / np.sqrt(rho2 + rho*s) / np.pi


def _pdf_alt(x, rho):
    # This version uses a reformulation that, for large x, reduces
    # the relative error and eliminates the overflow warning.
    #
    # For scalar x and rho only.
    #
    k = _k(rho)
    if x > 1e8*rho:
        x4 = np.power(x, -4)
        p = k*x4/(powm1(rho/x, 2)**2 + rho**2*x4)
    else:
        p = _pdf_v2(x, rho)
    return p


def rel_error(value, reference):
    if value == reference:
        return 0.0
    if reference == 0:
        return 1.0
    return abs((value - reference)/reference)


def check_pdf(rho=1, n=10000, start=1e-12, stop=1e50, dps=250):
    # Check the relative error of rel_breitwigner.pdf.
    x = np.geomspace(start, stop, n)
    sp_pdf = stats.rel_breitwigner.pdf(x, rho)
    v2_pdf = _pdf_v2(x, rho)
    alt_pdf = np.array([_pdf_alt(t, rho) for t in x])
    with mp.workdps(dps):
        mp_pdf = [rel_breitwigner.pdf(t, rho, 1) for t in x]
        sp_relerr = np.array([float(rel_error(p, mp))
                              for p, mp in zip(sp_pdf, mp_pdf)])
        v2_relerr = np.array([float(rel_error(p, mp))
                              for p, mp in zip(v2_pdf, mp_pdf)])
        alt_relerr = np.array([float(rel_error(p, mp))
                               for p, mp in zip(alt_pdf, mp_pdf)])
    return x, sp_relerr, v2_relerr, alt_relerr


if __name__ == "__main__":
    rho = 100
    print(f'{rho = }')
    x, relerr, relerr_v2, relerr_alt = check_pdf(rho=rho, n=60000,
                                                 start=1e-12, stop=1e20,
                                                 dps=200)
    alpha = 0.35
    plt.subplot(3, 1, 1)
    plt.plot(x, relerr, '.', alpha=alpha, label=f'{rho = } (scipy PR)',
             markersize=3)
    plt.semilogx()
    plt.grid(True)
    plt.title(f'PDF rel. error, {rho = }')
    plt.subplot(3, 1, 2)
    plt.plot(x, relerr_v2, '.', alpha=alpha, label=f'{rho = } (v2)',
             markersize=3)
    plt.semilogx()
    plt.grid(True)
    plt.subplot(3, 1, 3)
    plt.plot(x, relerr_alt, '.', alpha=alpha, label=f'{rho = } (alt)',
             markersize=3)
    plt.semilogx()
    plt.grid(True)

    plt.semilogx()
    # plt.grid()
    plt.xlabel('x')
    plt.tight_layout()
    # plt.ylabel('PDF relative error')
    # plt.legend(shadow=True, framealpha=1)
    plt.show()