diff --git a/.github/workflows/dev-pytest.yml b/.github/workflows/dev-pytest.yml index 6cb2961..aecac6f 100644 --- a/.github/workflows/dev-pytest.yml +++ b/.github/workflows/dev-pytest.yml @@ -1,7 +1,8 @@ name: Dev pytest (push) on: - push + push: + branches: [ "dev" ] jobs: build: diff --git a/.github/workflows/docs-build.yml b/.github/workflows/docs-build.yml index eb02fed..be7a41f 100644 --- a/.github/workflows/docs-build.yml +++ b/.github/workflows/docs-build.yml @@ -2,9 +2,9 @@ name: Documentation build (push main) on: push: - branches: [ "main" ] + branches: [ "main", "dev" ] pull_request: - branches: [ "main" ] + branches: [ "main", "dev" ] jobs: build: diff --git a/README.md b/README.md index cb3931a..2af6a83 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,8 @@ The package was developed in the course of doctoral studies at the [University o The package is divided in the following submodules: - The `constants` module includes a collection of physical and mathematical constants and related functions used throughout the package. -- The `utilities` module provides complex functions and statistical methods. +- The `utilities` module provides auxilary complex functions. +- The `data` module contains statistical methods. - The `kinematics` module accommodates functions necessary for particle kinematics such as phase space functions, etc. - The `dynamics` module houses integration vertices, particle propagators, channel definitions, etc. - The `integrals` module provides numerical integration routines for dispersion integrals. diff --git a/docs/source/bibliography.rst b/docs/source/bibliography.rst index 1c3ccce..9209998 100644 --- a/docs/source/bibliography.rst +++ b/docs/source/bibliography.rst @@ -6,9 +6,13 @@ } - Bibliography ============ .. bibliography:: - :style: plain \ No newline at end of file + :style: plain + + +.. .. raw:: html + +.. diff --git a/docs/source/coverage/report_styles.css b/docs/source/coverage/report_styles.css new file mode 100644 index 0000000..c2a64c5 --- /dev/null +++ b/docs/source/coverage/report_styles.css @@ -0,0 +1,7 @@ +p.pln { + font-size: 5pt; +} + +p.pln span.n { + display: none; +} \ No newline at end of file diff --git a/docs/source/data.rst b/docs/source/data.rst new file mode 100644 index 0000000..a364268 --- /dev/null +++ b/docs/source/data.rst @@ -0,0 +1,14 @@ +Data +==== + +.. automodule:: dispersionrelations.data + :special-members: __call__ + :members: + :undoc-members: + :show-inheritance: + +.. currentmodule:: dispersionrelations.data + +.. autosummary:: + :toctree: data_api + :recursive: \ No newline at end of file diff --git a/docs/source/literature.bib b/docs/source/literature.bib index ad962f0..e01f8d0 100644 --- a/docs/source/literature.bib +++ b/docs/source/literature.bib @@ -50,6 +50,19 @@ @article{Chung:1995dx year = "1995" } +@article{Schafer:2023qtl, + author = {Sch{\"a}fer, Hannah and Zanke, Marvin and Korte, Yannis and Kubis, Bastian}, + title = "{The semileptonic decays $\eta^{(\prime)} \to \pi^0 \ell^+ \ell^-$ and $\eta' \to \eta \ell^+ \ell^-$ in the standard model}", + eprint = "2307.10357", + archivePrefix = "arXiv", + primaryClass = "hep-ph", + doi = "10.1103/PhysRevD.108.074025", + journal = "Phys. Rev. D", + volume = "108", + number = "7", + pages = "074025", + year = "2023" +} @Article{ harris2020array, title = {Array programming with {NumPy}}, diff --git a/docs/source/modules.rst b/docs/source/modules.rst index 6b5c8fc..6afd124 100644 --- a/docs/source/modules.rst +++ b/docs/source/modules.rst @@ -8,13 +8,15 @@ The package consists of the following modules: constants utils + data kinematics dynamics integrals The **constants** module includes a collection of physical and mathematical constants and related functions used throughout the package. -The **utilities** module provides complex functions and statistical methods. +The **utilities** module provides auxilary complex functions. +The **data** module contains statistical methods. The **kinematics** module accommodates functions necessary for particle kinematics such as phase space functions, etc. (see :ref:`kinematics-section`). The **dynamics** module houses integration vertices, particle propagators, channel definitions, etc. (see :ref:`dynamics-section`). The **integrals** module provides numerical integration routines for dispersion integrals. \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 594422e..03022cc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "dispersionrelations" -version = "1.0.0" +version = "1.1.0" authors = [ {name="George Chanturia"} ] @@ -17,7 +17,7 @@ classifiers = [ "Programming Language :: Python :: 3", "Operating System :: OS Independent", ] -dependencies = ["numpy", "scipy", "scikit-image", "matplotlib"] +dependencies = ["numpy", "scipy", "scikit-image", "matplotlib", "uncertainties"] [project.urls] homepage = "https://github.com/giochanturia/DispersionRelations" @@ -33,7 +33,22 @@ docs = [ "sphinxcontrib-bibtex", "docutils", ] -dev = ["pytest"] +dev = ["pytest", "pytest-cov"] [tool.setuptools.packages.find] -where = ["src"] \ No newline at end of file +where = ["src"] + +[tool.pytest.ini_options] +addopts = "--cov=dispersionrelations --cov-report=html" + +[tool.coverage.run] +branch = true +source = ["dispersionrelations"] +omit = [ + "src/dispersionrelations/__init__.py", + "src/dispersionrelations/plotting.py", + "*/dispersionrelations/plotting.py", +] + +[tool.coverage.html] +extra_css = "docs/source/coverage/report_styles.css" \ No newline at end of file diff --git a/src/dispersionrelations/constants.py b/src/dispersionrelations/constants.py index 1dc4f3d..6a841d3 100644 --- a/src/dispersionrelations/constants.py +++ b/src/dispersionrelations/constants.py @@ -5,6 +5,7 @@ """ import numpy as np +from uncertainties import ufloat MeV = 1e-3 @@ -22,13 +23,18 @@ TeV = 1e6 * GeV r"""Teraelectronvolt, numerically defined as :math:`10^{6}\times` :code:`MeV`.""" -alpha_fs = 1 / 137 +E_almost_zero = 1e-5 * GeV +r"""A small energy, used in order to avoid division by zero.""" + +s_almost_zero = E_almost_zero**2 + +alpha_fs = 1 / 137.036 r"""Electromagnetic fine structure constant :math:`\alpha_{\text{fs}}\approx1/137`.""" e2 = 4 * np.pi * alpha_fs r"""Electric charge squared, :math:`e^2 = 4\pi\alpha_{\text{fs}}`.""" -to_nb = 3.894e5 * GeV**2 +to_nb = 1 / (2.56819 * 1e-6) * GeV**2 r"""Conversion factor from natural units to nanobarns. **Note**: a "natural unit" is defined by whatever is equal to `1` by default. @@ -47,10 +53,10 @@ M_PI0 = 134.98 * MeV r"""Mass of the neutral pion :cite:`Zyla:2020zbs`.""" -M_K = 493.68 * MeV +M_K = 493.677 * MeV r"""Mass of the charged kaon :cite:`Zyla:2020zbs`.""" -M_K0 = 497.61 * MeV +M_K0 = 497.611 * MeV r"""Mass of the neutral kaon :cite:`Zyla:2020zbs`.""" M_RHO = 762.5 * MeV @@ -122,7 +128,7 @@ G_KSTAR = 2 * 26 * MeV r"""Width of the :math:`K^*` :cite:`Zyla:2020zbs`.""" -G_KSTAR_STD = 6 * MeV +G_KSTAR_STD = 2 * 6 * MeV r"""Standard deviation of the :math:`K^*` width :cite:`Zyla:2020zbs`.""" M_A1_POLE = 1209 * MeV @@ -137,104 +143,69 @@ G_A1_POLE_STD = 100 * MeV # largest error r"""Standard deviation of the :math:`a_1` width, obtained from the pole position :cite:`JPAC:2018zwp`.""" +BR_OMEGA_TO_3PI = 89.2 / 100 +r"""Branching ratio of :math:`\omega \to 3\pi` decay :cite:`Zyla:2020zbs`.""" + +BR_OMEGA_TO_3PI_STD = 0.7 / 100 +r"""Standard deviation of the branching ratio of :math:`\omega \to 3\pi` decay :cite:`Zyla:2020zbs`.""" + +__G_OMEGA = ufloat(G_OMEGA, G_OMEGA_STD) +__BR_OMEGA_TO_3PI = ufloat(BR_OMEGA_TO_3PI, BR_OMEGA_TO_3PI_STD) +__G_OMEGA_TO_3PI = __G_OMEGA * __BR_OMEGA_TO_3PI + +G_OMEGA_TO_3PI = float(__G_OMEGA_TO_3PI.nominal_value) +r"""Partial width of :math:`\omega \to 3\pi` decay.""" + +G_OMEGA_TO_3PI_STD = float(__G_OMEGA_TO_3PI.std_dev) +r"""Standard deviation of the partial width of :math:`\omega \to 3\pi` decay.""" + +BR_OMEGA_TO_PI0GAMMA = 8.33 / 100 +r"""Branching ratio of :math:`\omega \to \pi^0 \gamma` decay :cite:`Zyla:2020zbs`.""" + +BR_OMEGA_TO_PI0GAMMA_STD = 0.25 / 100 +r"""Standard deviation of the branching ratio of :math:`\omega \to \pi^0 \gamma` decay :cite:`Zyla:2020zbs`.""" + +__BR_OMEGA_TO_PI0GAMMA = ufloat(BR_OMEGA_TO_PI0GAMMA, BR_OMEGA_TO_PI0GAMMA_STD) +__G_OMEGA_TO_PI0GAMMA = __G_OMEGA * __BR_OMEGA_TO_PI0GAMMA + +G_OMEGA_TO_PI0GAMMA = float(__G_OMEGA_TO_PI0GAMMA.nominal_value) +r"""Partial width of :math:`\omega \to \pi^0 \gamma` decay.""" + +G_OMEGA_TO_PI0GAMMA_STD = float(__G_OMEGA_TO_PI0GAMMA.std_dev) +r"""Standard deviation of the partial width of :math:`\omega \to \pi^0 \gamma` decay.""" + +G_RHO_TO_ETAGAMMA = 44.2 * keV +r"""Partial width of :math:`\rho \to \eta \gamma` decay :cite:`Schafer:2023qtl`.""" + +G_RHO_TO_ETAGAMMA_STD = 3.1 * keV +r"""Standard deviation of the partial width of :math:`\rho \to \eta \gamma` decay :cite:`Schafer:2023qtl`.""" + +BR_KSTAR_TO_KGAMMA_CHARGED = 0.098 / 100 +r"""Branching ratio of the charged :math:`K^* \to K \gamma` decay :cite:`Dax:2020dzg`.""" + +BR_KSTAR_TO_KGAMMA_CHARGED_STD = 0.009 / 100 +r"""Standard deviation of the branching ratio of the charged :math:`K^* \to K \gamma` decay :cite:`Dax:2020dzg`.""" + +BR_KSTAR_TO_KGAMMA_NEUTRAL = 0.246 / 100 +r"""Branching ratio of the neutral :math:`K^* \to K \gamma` decay :cite:`Dax:2020dzg`.""" + +BR_KSTAR_TO_KGAMMA_NEUTRAL_STD = 0.021 / 100 +r"""Standard deviation of the branching ratio of the neutral :math:`K^* \to K \gamma` decay :cite:`Dax:2020dzg`.""" + +__G_KSTAR = ufloat(G_KSTAR, G_KSTAR_STD) +__BR_KSTAR_TO_KGAMMA_CHARGED = ufloat(BR_KSTAR_TO_KGAMMA_CHARGED, BR_KSTAR_TO_KGAMMA_CHARGED_STD) +__BR_KSTAR_TO_KGAMMA_NEUTRAL = ufloat(BR_KSTAR_TO_KGAMMA_NEUTRAL, BR_KSTAR_TO_KGAMMA_NEUTRAL_STD) +__G_KSTAR_TO_KGAMMA_CHARGED = __G_KSTAR * __BR_KSTAR_TO_KGAMMA_CHARGED +__G_KSTAR_TO_KGAMMA_NEUTRAL = __G_KSTAR * __BR_KSTAR_TO_KGAMMA_NEUTRAL + +G_KSTAR_TO_KGAMMA_CHARGED = float(__G_KSTAR_TO_KGAMMA_CHARGED.nominal_value) +r"""Partial width of the charged :math:`K^* \to K \gamma` decay.""" + +G_KSTAR_TO_KGAMMA_CHARGED_STD = float(__G_KSTAR_TO_KGAMMA_CHARGED.std_dev) +r"""Standard deviation of the partial width of the charged :math:`K^* \to K \gamma` decay.""" + +G_KSTAR_TO_KGAMMA_NEUTRAL = float(__G_KSTAR_TO_KGAMMA_NEUTRAL.nominal_value) +r"""Partial width of the neutral :math:`K^* \to K \gamma` decay.""" -def scientific_notation(num, rounding=2): - r""" - Scientific notation. - - Parameters - ---------- - num : float - A real number. - rounding : int - Number of digits after the period. - - Returns - ------- - num_not : str - A LaTeX code for the number representation. - - Examples - -------- - >>> import numpy as np - >>> from DispersionRelations.constants import scientific_notation - >>> print(scientific_notation(np.pi)) - '3.1415926536' - >>> print(scientific_notation(13.4)) - '1.34 \\times 10^{1}' - """ - exponent = int(np.floor(np.log10(np.absolute(num)))) - radical = round(num / (10**exponent), rounding) - if exponent == 0: - return str(radical) - return str(radical) + r" \times 10^{" + str(exponent) + "}" - - -def rounding_PDG(mean, std): - r""" - Rounding with PDG rules :cite:`Zyla:2020zbs`. - - Parameters - ---------- - mean : float - Mean value of the quantity. - std : float - Standard deviation of the quantity. - - Returns - ------- - mean_out, std_out : (float, float) - Mean value and standard deviation, cited according to the PDG prescription. - - Examples - -------- - >>> from DispersionRelations.constants import rounding_PDG - >>> print(rounding_PDG(0.827, 0.119)) - (0.83, 0.12) - >>> print(rounding_PDG(0.827, 0.367)) - (0.8, 0.4) - """ - before_zero = int(np.floor(np.log10(np.absolute(std)))) + 1 - three_highest_digits = int(std / 10 ** (before_zero - 3)) - to_keep = 2 # default - if 100 <= three_highest_digits <= 354: - to_keep = 2 - if 355 <= three_highest_digits <= 949: - to_keep = 1 - if 950 <= three_highest_digits <= 999: - to_keep = 1 # because 1000 has 3 digits now - three_highest_digits = 1000 - - std_out = round( - three_highest_digits * 10 ** (before_zero - 3), to_keep - before_zero - ) - mean_out = round(mean, to_keep - before_zero) - - if mean_out % 1 == 0 and std_out % 1 == 0: - mean_out = int(mean_out) - std_out = int(std_out) - - return mean_out, std_out - - -def sR(M_R, G_R): - r""" - Resonance pole location calculated from its mass and width. - - Parameters - ---------- - M_R : float - Mass of the particle. - G_R : float - Width of the particle. - - Returns - ------- - sR : complex - The complex pole location. - - Notes - ----- - The pole location is calculated using :math:`s_R(M_R, \Gamma_R) = (M_R - i \Gamma_R / 2)^2`. - """ - return (M_R - 1j * G_R / 2) ** 2 +G_KSTAR_TO_KGAMMA_NEUTRAL_STD = float(__G_KSTAR_TO_KGAMMA_NEUTRAL.std_dev) +r"""Standard deviation of the partial width of the neutral :math:`K^* \to K \gamma` decay.""" \ No newline at end of file diff --git a/src/dispersionrelations/data.py b/src/dispersionrelations/data.py new file mode 100644 index 0000000..242d669 --- /dev/null +++ b/src/dispersionrelations/data.py @@ -0,0 +1,1289 @@ +import numpy as np +import scipy as sp +import warnings +import os + +from dispersionrelations.utils import cite, uncite + + +def check_pd(mat): + r""" + Check if a matrix is positive definite. + + Parameters + ---------- + mat : array-like, shape (N, N) + The matrix to be checked. + + Returns + ------- + bool + True if the matrix is positive definite, False otherwise. + """ + try: + np.linalg.cholesky(mat) + return True + except np.linalg.LinAlgError: + return False + + +def corr_to_cov(corr, std): + r""" + Convert a correlation matrix to a covariance matrix. + + Parameters + ---------- + corr : array-like, shape (N, N) + The correlation matrix to be converted. + std : array-like, shape (N,) + The standard deviations corresponding to the variables. + + Returns + ------- + cov : array-like, shape (N, N) + The resulting covariance matrix. + + Raises + ------ + ValueError + If corr is not a square matrix or if corr and std have incompatible dimensions. + """ + corr = np.asarray(corr) + std = np.asarray(std) + if corr.shape[0] != corr.shape[1]: + raise ValueError(f"corr must be a square matrix, but got shape {corr.shape}") + if corr.shape[0] != std.shape[0]: + raise ValueError( + f"corr and std must have compatible dimensions, but got {corr.shape} and {std.shape}" + ) + if std.ndim != 1: + raise ValueError( + f"std must be a one-dimensional array, but got shape {std.shape}" + ) + outer = np.outer(std, std) + cov = corr * outer + return cov + + +def cov_to_corr(cov): + r""" + Convert a covariance matrix to a correlation matrix. + + Parameters + ---------- + cov : array-like, shape (N, N) + The covariance matrix to be converted. + + Returns + ------- + corr : array-like, shape (N, N) + The resulting correlation matrix. + + Raises + ------ + ValueError + If cov is not a square matrix. + """ + cov = np.asarray(cov) + if cov.shape[0] != cov.shape[1]: + raise ValueError(f"cov must be a square matrix, but got shape {cov.shape}") + corr = np.zeros(cov.shape) + std = np.sqrt(np.diag(cov)) + outer = np.outer(std, std) + nonzero = np.where(outer != 0) + corr[nonzero] = cov[nonzero] / outer[nonzero] + corr[np.isnan(corr)] = 0 + return corr + + +def cov_to_std(cov): + r""" + Convert a covariance matrix to standard deviations (square root of the diagonal elements). + + Parameters + ---------- + cov : array-like, shape (N, N) + The covariance matrix to be converted. + + Returns + ------- + std : array-like, shape (N,) + The resulting standard deviations. + + Raises + ------ + ValueError + If cov is not a square matrix. + """ + cov = np.asarray(cov) + if cov.shape[0] != cov.shape[1]: + raise ValueError(f"cov must be a square matrix, but got shape {cov.shape}") + return np.sqrt(np.diag(cov)) + + +def std_to_cov(std, correlation=0.0): + r""" + Convert standard deviations to a covariance matrix, assuming a given correlation structure. + + Parameters + ---------- + std : array-like, shape (N,) + The standard deviations to be converted. + correlation : float, optional + The correlation coefficient to be used in the conversion. Default is 0.0 (no correlation). + + Returns + ------- + cov : array-like, shape (N, N) + The resulting covariance matrix. + + Raises + ------ + ValueError + If std is not a one-dimensional array. + """ + std = np.asarray(std) + if std.ndim != 1: + raise ValueError( + f"std must be a one-dimensional array, but got shape {std.shape}" + ) + corr = (1.0 - correlation) * np.eye(len(std)) + correlation * np.ones( + (len(std), len(std)) + ) + cov = np.outer(std, std) * corr + return cov + + +def chi2(fit_data, observed_data, covariance_matrix, weights=None): + r""" + :math:`\chi^2`, computed using modelled and observed data. + + Parameters + ---------- + fit_data : array_like + Data obtained from a model, :math:`f_i \in \{f_1, f_2, \dots, f_n\}`. + observed_data : array_like + Data obtained from an experiment, :math:`o_i \in \{o_1, o_2, \dots, o_n\}`. + covariance_matrix : array_like + Covariance matrix of the experimental data, :math:`C \in \mathbb{R}^{n\times n}`. + weights : array_like + Weight matrix, used to manually attenuate parts of the data, :math:`W \in \mathbb{R}^{n\times n}`. + + Returns + ------- + s : float + The total :math:`\chi^2`. + + Notes + ----- + The :math:`\chi^2` is calculated using + + .. math:: \chi^2 = \sum_i \sum_j (f_i-o_i)(C^{-1}_{ij}W_{ij})(f_j-o_j). + + See also + -------- + dispersionrelations.data.chi2_vector : for individual contributions. + """ + if weights is not None: + covariance_matrix = covariance_matrix / weights + difference = fit_data - observed_data + return difference @ np.linalg.inv(covariance_matrix) @ difference + + +def chi2_with_inverse(fit_data, observed_data, inverse_covariance_matrix, weights=None): + r""" + :math:`\chi^2`, computed using modelled and observed data. + This version of the function accepts the inverse of the covariance matrix + as a parameter and therefore gains speed by avoiding matrix inversion. + + Parameters + ---------- + fit_data : array_like + Data obtained from a model, :math:`f_i \in \{f_1, f_2, \dots, f_n\}`. + observed_data : array_like + Data obtained from an experiment, :math:`o_i \in \{o_1, o_2, \dots, o_n\}`. + inverse_covariance_matrix : array_like + Inverted covariance matrix of the experimental data, :math:`I \in \mathbb{R}^{n\times n}`. + weights : array_like + Weight matrix, used to manually attenuate parts of the data, :math:`W \in \mathbb{R}^{n\times n}`. + + Returns + ------- + s : float + The total :math:`\chi^2`. + + Notes + ----- + The :math:`\chi^2` is calculated using + + .. math:: \chi^2 = \sum_i \sum_j (f_i-o_i)(I_{ij}W_{ij})(f_j-o_j). + + See also + -------- + dispersionrelations.data.chi2_vector_with_inverse : for individual contributions. + """ + if weights is not None: + inverse_covariance_matrix = inverse_covariance_matrix * weights + difference = fit_data - observed_data + return difference @ inverse_covariance_matrix @ difference + + +def chi2_vector(fit_data, observed_data, covariance_matrix, weights=None): + r""" + :math:`\chi^2` vector. + + Parameters + ---------- + fit_data : array_like + Data obtained from a model, :math:`f_i \in \{f_1, f_2, \dots, f_n\}`. + observed_data : array_like + Data obtained from an experiment, :math:`o_i \in \{o_1, o_2, \dots, o_n\}`. + covariance_matrix : array_like + Covariance matrix of the experimental data, :math:`C \in \mathbb{R}^{n\times n}`. + weights : array_like + Weight matrix, used to manually attenuate parts of the data, :math:`W \in \mathbb{R}^{n\times n}`. + + Returns + ------- + v : array_like + A vector with shape (n,) of individual :math:`\chi^2` contributions. + + Notes + ----- + The :math:`\chi^2` vector is built using + + .. math:: v_i = \sum_j (f_i-o_i)(C^{-1}_{ij}W_{ij})(f_j-o_j). + + Warning + ------- + When the covariance matrix is non-diagonal, the individual components of this + vector do not have a statistical meaning, since the observations are correlated. + Use with caution! + + Examples + -------- + >>> import numpy as np + >>> from dispersionrelations.data import chi2_vector + >>> np.random.seed(137) + >>> n = 10 + >>> x_f = np.linspace(1, 2, n) + >>> x_o = x_f + 0.1 * np.random.rand(n) + >>> x_e = 0.1 * (np.random.rand(n) + 1) / 2 + >>> x_C = np.diag(x_e)**2 + >>> print(chi2_vector(x_f, x_o, x_C)) + [0.89280398 0.01007333 0.62584069 0.67718574 0.32026806 0.99981995 0.91740446 0.01387227 2.33635045 0.18841755] + + .. plot:: + + import numpy as np + from dispersionrelations.data import chi2_vector + import matplotlib.pyplot as plt + + np.random.seed(137) + n = 10 + x_f = np.linspace(1, 2, n) + x_o = x_f + 0.1 * np.random.rand(n) + x_e = 0.1 * (np.random.rand(n) + 1) / 2 + x_C = np.diag(x_e)**2 + x_chi2_vector = chi2_vector(x_f, x_o, x_C) + plt.errorbar(np.arange(n), x_o, yerr=x_e, fmt="o", markersize=0, capsize=2, label="observed_data", zorder=2) + plt.plot(np.arange(n), x_f, label="fit_data", zorder=2) + plt.fill_between(np.arange(n), 0.1 * x_chi2_vector, 0, color='tab:red', label=r"$\chi^2/10$", zorder=3) + plt.legend() + + See also + -------- + dispersionrelations.data.chi2 + """ + if weights is not None: + covariance_matrix = covariance_matrix / weights + difference = fit_data - observed_data + return difference * (np.linalg.inv(covariance_matrix) @ difference) + + +def chi2_vector_with_inverse( + fit_data, observed_data, inverse_covariance_matrix, weights=None +): + r""" + :math:`\chi^2`, computed using modelled and observed data. + This version of the function accepts the inverse of the covariance matrix + as a parameter and therefore gains speed by avoiding matrix inversion. + + Parameters + ---------- + fit_data : array_like + Data obtained from a model, :math:`f_i \in \{f_1, f_2, \dots, f_n\}`. + observed_data : array_like + Data obtained from an experiment, :math:`o_i \in \{o_1, o_2, \dots, o_n\}`. + inverse_covariance_matrix : array_like + Inverted covariance matrix of the experimental data, :math:`I \in \mathbb{R}^{n\times n}`. + weights : array_like + Weight matrix, used to manually attenuate parts of the data, :math:`W \in \mathbb{R}^{n\times n}`. + + Returns + ------- + v : array_like + A vector with shape (n,) of individual :math:`\chi^2` contributions. + + Notes + ----- + The :math:`\chi^2` vector is built using + + .. math:: v_i = \sum_j (f_i-o_i)(I_{ij}W_{ij})(f_j-o_j). + + See also + -------- + dispersionrelations.data.chi2_with_inverse + """ + if weights is not None: + inverse_covariance_matrix = inverse_covariance_matrix * weights + difference = fit_data - observed_data + return difference * (inverse_covariance_matrix @ difference) + + +def chi2_reduced( + fit_data, observed_data, covariance_matrix, weights=None, number_of_parameters=0 +): + r""" + A reduced :math:`\chi^2`, computed using modelled and observed data. + + Parameters + ---------- + fit_data : array_like + Data obtained from a model, :math:`f_i \in \{f_1, f_2, \dots, f_n\}`. + observed_data : array_like + Data obtained from an experiment, :math:`o_i \in \{o_1, o_2, \dots, o_n\}`. + covariance_matrix : array_like + Covariance matrix of the experimental data, :math:`C \in \mathbb{R}^{n\times n}`. + weights : array_like + Weight matrix, used to manually attenuate parts of the data, :math:`W \in \mathbb{R}^{n\times n}`. + number_of_parameters : int + Number of parameters in the model, used for calculating the degrees of freedom. + + Returns + ------- + s : float + The reduced :math:`\chi^2`. + + Notes + ----- + The reduced :math:`\chi^2` is calculated using + + .. math:: \bar{\chi}^2 = \frac{\chi^2}{\text{d.o.f.}} = \frac{1}{n - n_\text{par}} \sum_i \sum_j (f_i-o_i)(C^{-1}_{ij}W_{ij})(f_j-o_j). + + See also + -------- + dispersionrelations.data.chi2 + """ + degrees_of_freedom = len(observed_data) - number_of_parameters + return ( + chi2(fit_data, observed_data, covariance_matrix, weights) / degrees_of_freedom + ) + + +def chi2_reduced_with_inverse( + fit_data, + observed_data, + inverse_covariance_matrix, + weights=None, + number_of_parameters=0, +): + r""" + A reduced :math:`\chi^2`, computed using modelled and observed data. + This version of the function accepts the inverse of the covariance matrix + as a parameter and therefore gains speed by avoiding matrix inversion. + + Parameters + ---------- + fit_data : array_like + Data obtained from a model, :math:`f_i \in \{f_1, f_2, \dots, f_n\}`. + observed_data : array_like + Data obtained from an experiment, :math:`o_i \in \{o_1, o_2, \dots, o_n\}`. + inverse_covariance_matrix : array_like + Inverted covariance matrix of the experimental data, :math:`I \in \mathbb{R}^{n\times n}`. + weights : array_like + Weight matrix, used to manually attenuate parts of the data, :math:`W \in \mathbb{R}^{n\times n}`. + number_of_parameters : int + Number of parameters in the model, used for calculating the degrees of freedom. + + Returns + ------- + s : float + The reduced :math:`\chi^2`. + + Notes + ----- + The reduced :math:`\chi^2` is calculated using + + .. math:: \bar{\chi}^2 = \frac{\chi^2}{\text{d.o.f.}} = \frac{1}{n - n_\text{par}} \sum_i \sum_j (f_i-o_i)(I_{ij}W_{ij})(f_j-o_j). + + See also + -------- + dispersionrelations.data.chi2 + """ + degrees_of_freedom = len(observed_data) - number_of_parameters + return ( + chi2_with_inverse(fit_data, observed_data, inverse_covariance_matrix, weights) + / degrees_of_freedom + ) + + +class Datablock: + r""" + A class to represent a dataset with x and y values, along with statistical and systematic covariance matrices. + + Parameters + ---------- + x : array-like, shape (N,) + The x values of the data points. + y : array-like, shape (N,) + The y values of the data points. + cov_stat : array-like, shape (N, N), optional + The initial statistical covariance matrix. Default is None, which initializes it to a zero matrix. + cov_syst : array-like, shape (N, N), optional + The initial systematic covariance matrix. Default is None, which initializes it to a zero matrix. + metadata : dict, optional + A dictionary to store additional metadata about the dataset. Default is an empty dictionary. + Common keys include 'ref', 'header', 'footer', utilised by the :class:`Datawriter` class. + + Attributes + ---------- + x : array-like, shape (N,) + The x values of the data points. + y : array-like, shape (N,) + The y values of the data points. + N : int + The number of data points. + cov_stat : array-like, shape (N, N) + The statistical covariance matrix. + cov_syst : array-like, shape (N, N) + The systematic covariance matrix. + metadata : dict + A dictionary to store additional metadata about the dataset. + """ + + def __init__(self, x, y, xlabel, ylabel, cov_stat=None, cov_syst=None, metadata={}): + self.x = x + self.y = y + self.xlabel = xlabel + self.ylabel = ylabel + self.N = len(x) + if len(y) != self.N: + raise ValueError( + f"x and y must have the same length, but got {len(x)} and {len(y)}" + ) + self.cov_stat = np.zeros((self.N, self.N)) + self.cov_syst = np.zeros((self.N, self.N)) + + if cov_stat is not None: + self.add_cov_stat(cov_stat) + self.check_cov_pd() + if cov_syst is not None: + self.add_cov_syst(cov_syst) + self.check_cov_pd() + + self.metadata = metadata + + def check_cov_pd(self): + r""" + Check if the total covariance matrix (statistical + systematic) is positive definite. + Display a warning if it is not. + + Returns + ------- + None + """ + if check_pd(self.cov_stat + self.cov_syst) == False: + warnings.warn("The total covariance matrix is not positive definite") + + def get_std_stat(self): + r""" + Get the statistical standard deviations (square root of the diagonal elements of the statistical covariance matrix). + + Returns + ------- + std_stat : array-like, shape (N,) + The statistical standard deviations. + """ + return cov_to_std(self.cov_stat) + + def get_std_syst(self): + r""" + Get the systematic standard deviations (square root of the diagonal elements of the systematic covariance matrix). + + Returns + ------- + std_syst : array-like, shape (N,) + The systematic standard deviations. + """ + return cov_to_std(self.cov_syst) + + def get_std_total(self, add_in_quadrature=True): + r""" + Get the total standard deviations (square root of the diagonal elements of the total covariance matrix). + + Parameters + ---------- + add_in_quadrature : bool, optional + If True, compute the total standard deviations by adding the statistical and systematic standard deviations in quadrature. + If False, compute the total standard deviations from the total covariance matrix. Default is True. + + Returns + ------- + std_total : array-like, shape (N,) + The total standard deviations. + """ + std_stat = self.get_std_stat() + std_syst = self.get_std_syst() + if add_in_quadrature: + return np.sqrt(std_stat**2 + std_syst**2) + else: + return std_stat + std_syst + + def add_cov_stat(self, cov_stat): + r""" + Add a contribution to the statistical covariance matrix. + + Parameters + ---------- + cov_stat : array-like, shape (N, N) + The statistical covariance matrix to be added. + + Returns + ------- + None + + Raises + ------ + ValueError + If cov_stat is not of shape (N, N). + """ + if cov_stat.shape != (self.N, self.N): + raise ValueError( + f"cov_stat must be of shape (N, N), but got {cov_stat.shape}" + ) + self.cov_stat += cov_stat + + def add_cov_syst(self, cov_syst): + r""" + Add a contribution to the systematic covariance matrix. + + Parameters + ---------- + cov_syst : array-like, shape (N, N) + The systematic covariance matrix to be added. + + Returns + ------- + None + + Raises + ------ + ValueError + If cov_syst is not of shape (N, N). + """ + if cov_syst.shape != (self.N, self.N): + raise ValueError( + f"cov_syst must be of shape (N, N), but got {cov_syst.shape}" + ) + self.cov_syst += cov_syst + + def add_std_stat(self, std_stat, correlation=0.0): + r""" + Add a contribution to the statistical covariance matrix from standard deviations. + + Parameters + ---------- + std_stat : array-like, shape (N,) or scalar + The standard deviations to be converted to covariance matrix. + correlation : float, optional + The correlation coefficient to be used in the conversion. Default is 0.0 (no correlation). + + Returns + ------- + None + + Raises + ------ + ValueError + If std_stat is not a scalar or of shape (N,). + """ + std_stat = np.asarray(std_stat) + if std_stat.shape != (self.N,) and std_stat.shape != (1,): + raise ValueError( + f"std_stat must be a scalar or of shape (N,), but got {std_stat.shape}" + ) + cov_stat = std_to_cov(std_stat, correlation=correlation) + self.add_cov_stat(cov_stat) + + def add_rel_stat(self, rel_stat, correlation=0.0): + r""" + Add a contribution to the statistical covariance matrix from relative uncertainties. + + Parameters + ---------- + rel_stat : array-like, shape (N,) or scalar + The relative uncertainties to be converted to standard deviations. + correlation : float, optional + The correlation coefficient to be used in the conversion. Default is 0.0 (no correlation). + + Returns + ------- + None + + Raises + ------ + ValueError + If rel_stat is not a scalar or of shape (N,). + """ + std_stat = np.abs(rel_stat * self.y) + self.add_std_stat(std_stat, correlation=correlation) + + def add_std_syst(self, std_syst, correlation=1.0): + r""" + Add a contribution to the systematic covariance matrix from standard deviations. + + Parameters + ---------- + std_syst : array-like, shape (N,) or scalar + The standard deviations to be converted to covariance matrix. + correlation : float, optional + The correlation coefficient to be used in the conversion. Default is 1.0 (full correlation). + + Returns + ------- + None + + Raises + ------ + ValueError + If std_syst is not a scalar or of shape (N,). + """ + std_syst = np.asarray(std_syst) + if std_syst.shape != (self.N,) and std_syst.shape != (1,): + raise ValueError( + f"std_syst must be a scalar or of shape (N,), but got {std_syst.shape}" + ) + cov_syst = std_to_cov(std_syst, correlation=correlation) + self.add_cov_syst(cov_syst) + + def add_rel_syst(self, rel_syst, correlation=1.0): + r""" + Add a contribution to the systematic covariance matrix from relative uncertainties. + + Parameters + ---------- + rel_syst : array-like, shape (N,) or scalar + The relative uncertainties to be converted to standard deviations. + correlation : float, optional + The correlation coefficient to be used in the conversion. Default is 1.0 (full correlation). + + Returns + ------- + None + + Raises + ------ + ValueError + If rel_syst is not a scalar or of shape (N,). + """ + std_syst = np.abs(rel_syst * self.y) + self.add_std_syst(std_syst, correlation=correlation) + + def add_external_calibration(self, c, c_std, correlation=1.0, new_ylabel=None): + r""" + Apply an external calibration factor to the data and propagate its uncertainty. + + Parameters + ---------- + c : float + The calibration factor to be applied to the data. + c_std : float + The uncertainty of the calibration factor. + correlation : float, optional + The correlation coefficient to be used in the propagation of the calibration uncertainty. Default is 1.0 (full correlation). + new_ylabel : str, optional + The new label for the y-axis after calibration. If None, the label remains unchanged. + + Returns + ------- + None + + Raises + ------ + ValueError + If c or c_std is not a scalar. + + Notes + ----- + For x-dependent calibration factors, use `rescale_y` instead. + """ + if np.asarray(c).shape != () or np.asarray(c_std).shape != (): + raise ValueError("c and c_std must be scalars") + self.y = c * self.y + self.cov_stat *= c**2 + self.cov_syst *= c**2 + if new_ylabel is not None: + self.ylabel = new_ylabel + if c_std != 0: + self.add_rel_syst(c_std / c, correlation=correlation) + self.check_cov_pd() + + def rescale_y( + self, + factor, + factor_std=None, + stat_or_syst="syst", + correlation=0.0, + new_ylabel=None, + ): + r""" + Rescale the y values by a given factor without propagating the uncertainty. + + Parameters + ---------- + factor : array-like, shape (N,) + The factor by which to rescale the y values. If an array is provided, it must be of shape (N,). + factor_std : array-like, shape (N,) + The uncertainty of the rescaling factor (not propagated). + stat_or_syst : str, optional + Specify whether the uncertainty is 'stat' (statistical) or 'syst' (systematic). Default is 'syst'. + correlation : float, optional + The correlation coefficient to be used for the uncertainty propagation. Default is 0.0 (no correlation). + new_ylabel : str, optional + The new label for the y-axis after rescaling. If None, the label remains unchanged + + Returns + ------- + None + + Raises + ------ + ValueError + If factor is not a scalar or of shape (N,). + + Notes + ----- + For a global rescaling, use `add_external_calibration` instead. + """ + factor = np.asarray(factor) + factor_std = ( + np.asarray(factor_std) if factor_std is not None else np.zeros(factor.shape) + ) + if factor.shape != (self.N,): + raise ValueError( + f"factor must be a scalar or of shape (N,), but got {factor.shape}" + ) + if factor_std.shape != (self.N,): + raise ValueError( + f"factor_std must be a scalar or of shape (N,), but got {factor_std.shape}" + ) + if stat_or_syst not in ["stat", "syst"]: + raise ValueError("stat_or_syst must be 'stat' or 'syst'") + self.y = factor * self.y + factor_mat = np.outer(factor, factor) + self.cov_stat = factor_mat * self.cov_stat + self.cov_syst = factor_mat * self.cov_syst + if new_ylabel is not None: + self.ylabel = new_ylabel + if np.any(factor_std != 0): + rel_syst = np.abs(factor_std / factor) + if stat_or_syst == "stat": + self.add_rel_stat(rel_syst, correlation=correlation) + elif stat_or_syst == "syst": + self.add_rel_syst(rel_syst, correlation=correlation) + self.check_cov_pd() + + def filter_data(self, mask): + r""" + Filter the data based on a boolean mask. + + Parameters + ---------- + mask : array-like, shape (N,) + A boolean array indicating which data points to keep. + + Returns + ------- + None + + Raises + ------ + ValueError + If mask is not of length N. + """ + mask = np.asarray(mask) + if mask.shape != (self.N,): + raise ValueError( + f"The mask must be of length N={self.N}, but got {mask.shape}" + ) + self.x = self.x[mask] + self.y = self.y[mask] + self.cov_stat = self.cov_stat[np.ix_(mask, mask)] + self.cov_syst = self.cov_syst[np.ix_(mask, mask)] + self.N = len(self.x) + self.check_cov_pd() + + def filter_xrange(self, x_min=None, x_max=None): + r""" + Filter the data based on x range. + + Parameters + ---------- + x_min : float, optional + The minimum x value to keep. If None, no lower bound is applied. + x_max : float, optional + The maximum x value to keep. If None, no upper bound is applied. + + Returns + ------- + None + """ + mask_lower = ( + self.x >= x_min if x_min is not None else np.ones(self.N, dtype=bool) + ) + mask_upper = ( + self.x <= x_max if x_max is not None else np.ones(self.N, dtype=bool) + ) + self.filter_data(mask_lower * mask_upper) + + def filter_yrange(self, y_min=None, y_max=None): + r""" + Filter the data based on y range. + + Parameters + ---------- + y_min : float, optional + The minimum y value to keep. If None, no lower bound is applied. + y_max : float, optional + The maximum y value to keep. If None, no upper bound is applied. + + Returns + ------- + None + """ + mask_lower = ( + self.y >= y_min if y_min is not None else np.ones(self.N, dtype=bool) + ) + mask_upper = ( + self.y <= y_max if y_max is not None else np.ones(self.N, dtype=bool) + ) + self.filter_data(mask_lower * mask_upper) + + def remove_datapoints(self, indices): + r""" + Remove data points based on their indices. + + Parameters + ---------- + indices : array-like, shape (M,) + An array of indices of the data points to be removed. + + Returns + ------- + None + + Raises + ------ + ValueError + If any index is out of bounds. + """ + indices = np.asarray(indices) + if np.any(indices < 0) or np.any(indices >= self.N): + raise ValueError( + f"All indices must be in the range [0, {self.N}), but got {indices}" + ) + mask = np.ones(self.N, dtype=bool) + mask[indices] = False + self.filter_data(mask) + + def downsample_by(self, factor): + r""" + Downsample the data by a given factor. + + Parameters + ---------- + factor : int + The downsampling factor. Must be a positive integer. + + Returns + ------- + None + + Raises + ------ + ValueError + If factor is not a positive integer. + """ + if not isinstance(factor, int) or factor <= 0: + raise ValueError("factor must be a positive integer") + mask = np.zeros(self.N, dtype=bool) + mask[::factor] = True + self.filter_data(mask) + + def downsample_to(self, N_points): + r""" + Downsample the data to a given number of points. + + Parameters + ---------- + N_points : int + The desired number of data points after downsampling. Must be a positive integer less than or equal to N. + + Returns + ------- + None + + Raises + ------ + ValueError + If N_points is not a positive integer or greater than N. + """ + raise NotImplementedError("downsample_to is not implemented yet") + + def interpolate(self, new_x): + r""" + Interpolate the data to new x values using linear interpolation. + + Parameters + ---------- + new_x : array-like, shape (M,) + The new x values to interpolate to. + + Returns + ------- + None + + Raises + ------ + ValueError + If new_x is not a one-dimensional array. + """ + raise NotImplementedError("interpolate is not implemented yet") + + def __str__(self): + result = f"Datablock with {self.N} points(s): {self.xlabel} vs {self.ylabel}." + if "ref" in self.metadata: + result += f" Reference: {self.metadata['ref']}." + return result + + def __repr__(self): + return self.__str__() + + def __add__(self, other): + r""" + Add two Datablock objects. + + Parameters + ---------- + other : Datablock + The other Datablock to be added. + + Returns + ------- + Datablock + A new Datablock containing the combined data. + + Raises + ------ + ValueError + If other is not a Datablock. + + Notes + ----- + Metadata and labels are not fully preserved in the addition. + """ + if not isinstance(other, Datablock): + raise ValueError("Can only add Datablock to Datablock") + + new_metadata = {} + if "ref" in self.metadata: + new_metadata["ref"] = self.metadata["ref"] + if "ref" in other.metadata: + if "ref" in new_metadata: + new_metadata["ref"] += ", " + other.metadata["ref"] + else: + new_metadata["ref"] = other.metadata["ref"] + warnings.warn("Metadata and labels are not preserved in Datablock addition") + + new_x = np.concatenate((self.x, other.x)) + new_y = np.concatenate((self.y, other.y)) + new_cov_stat = sp.linalg.block_diag(self.cov_stat, other.cov_stat) + new_cov_syst = sp.linalg.block_diag(self.cov_syst, other.cov_syst) + return Datablock( + new_x, + new_y, + "x", + "y", + cov_stat=new_cov_stat, + cov_syst=new_cov_syst, + metadata=new_metadata, + ) + + def __radd__(self, other): + if other == 0: + return self + return self.__add__(other) + + +class Datawriter: + r""" + A class for writing Datablock data to files. + + Parameters + ---------- + datablock : Datablock + The Datablock object containing the data to be written. + file_prefix : str + The prefix for the output files. + folder_path : str, optional + The folder path where the files will be saved. Default is the current directory. + If the folder does not exist, it will be created. + quiet : bool, optional + If True, suppresses output messages. Default is False. + format : dict, optional + A dictionary specifying the format for saving the data. + """ + + def __init__( + self, + datablock, + file_prefix, + folder_path="", + quiet=False, + format=dict(fmt="% .11e", delimiter="\t"), + ): + self.datablock = datablock + self.folder = folder_path + self.file_prefix = file_prefix + self.header = "" + self.footer = "" + self.quiet = quiet + self.format = format + if "header" in self.datablock.metadata: + self.header += self.datablock.metadata["header"] + if "footer" in self.datablock.metadata: + self.footer += self.datablock.metadata["footer"] + os.makedirs(self.folder, exist_ok=True) + + def write_data(self, **filekwargs): + r""" + Write the data to a .dat file. + The file will contain columns for x, y, statistical uncertainty, and systematic uncertainty. + + Parameters + ---------- + filekwargs : dict + Additional keyword arguments to be passed to np.savetxt. + + Returns + ------- + None + """ + data = np.column_stack( + ( + self.datablock.x, + self.datablock.y, + np.sqrt(np.diag(self.datablock.cov_stat)), + np.sqrt(np.diag(self.datablock.cov_syst)), + ) + ) + np.savetxt( + os.path.join(self.folder, f"{self.file_prefix}.dat"), + data, + header=self.header + + "\n" + + f"Columns {data.shape}: {self.datablock.xlabel}, {self.datablock.ylabel}, stat_unc, syst_unc", + footer=self.footer, + **self.format, + **filekwargs, + ) + if not self.quiet: + print( + f'"{os.path.join(self.folder, f"{self.file_prefix}.dat")}" created {data.shape}.' + ) + + def write_cov_stat(self, **filekwargs): + r""" + Write the statistical covariance matrix to a .dat file. + The file will contain the statistical covariance matrix. + + Parameters + ---------- + filekwargs : dict + Additional keyword arguments to be passed to np.savetxt. + + Returns + ------- + None + """ + np.savetxt( + os.path.join(self.folder, f"{self.file_prefix}_COV_STAT.dat"), + self.datablock.cov_stat, + header=self.header + + "\n" + + f"Statistical covariance matrix {self.datablock.cov_stat.shape}", + footer=self.footer, + **self.format, + **filekwargs, + ) + if not self.quiet: + print( + f'"{os.path.join(self.folder, f"{self.file_prefix}_COV_STAT.dat")}" created {self.datablock.cov_stat.shape}.' + ) + + def write_cov_syst(self, **filekwargs): + r""" + Write the systematic covariance matrix to a .dat file. + The file will contain the systematic covariance matrix. + + Parameters + ---------- + filekwargs : dict + Additional keyword arguments to be passed to np.savetxt. + + Returns + ------- + None + """ + np.savetxt( + os.path.join(self.folder, f"{self.file_prefix}_COV_SYST.dat"), + self.datablock.cov_syst, + header=self.header + + "\n" + + f"Systematic covariance matrix {self.datablock.cov_syst.shape}", + footer=self.footer, + **self.format, + **filekwargs, + ) + if not self.quiet: + print( + f'"{os.path.join(self.folder, f"{self.file_prefix}_COV_SYST.dat")}" created {self.datablock.cov_syst.shape}.' + ) + + def write_all(self, **filekwargs): + r""" + Write all data to .dat files. + This includes the main data file, the statistical covariance matrix file, and the systematic covariance matrix file. + + Parameters + ---------- + filekwargs : dict + Additional keyword arguments to be passed to np.savetxt. + This argument is shared between all three writing functions. + For a finer control, call the individual writing functions instead. + + Returns + ------- + None + """ + if not self.quiet: + print(f"- {self.file_prefix} - - - - - - - - - - - - - -") + self.write_data(**filekwargs) + self.write_cov_stat(**filekwargs) + self.write_cov_syst(**filekwargs) + + +class Datareader: + r""" + A class for reading Datablock data from files. + + Parameters + ---------- + folder : str + The folder path where the files are located. + prefix : str + The prefix for the input files. + """ + + def __init__(self, folder, prefix): + self.folder = folder + self.prefix = prefix + self.file_data = os.path.join(folder, f"{prefix}.dat") + self.file_cov_stat = os.path.join(folder, f"{prefix}_COV_STAT.dat") + self.file_cov_syst = os.path.join(folder, f"{prefix}_COV_SYST.dat") + + def read_all(self, **filekwargs): + r""" + Read all data from .dat files and return a Datablock object. + This includes the main data file, the statistical covariance matrix file, and the systematic covariance matrix file. + + Parameters + ---------- + filekwargs : dict + Additional keyword arguments to be passed to np.loadtxt. + + Returns + ------- + Datablock + A Datablock object containing the data read from the files. + + Raises + ------ + ValueError + If the shapes of the covariance matrices do not match the expected dimensions. + If the uncertainties (squared) do not match the diagonals of the covariance matrices. + """ + # Load data from file + data = np.loadtxt(self.file_data, **filekwargs) + + # Load comments from file + comments = [] + with open(self.file_data, encoding="utf-8") as f: + for line in f: + if line.lstrip().startswith("#"): + comments.append(line.lstrip()[1:].lstrip().rstrip("\n")) + + # Fill labels and metadata from comments + xlabel = "x" + ylabel = "y" + metadata = {} + try: + metadata["ref"] = uncite(comments[0]) + except: + warnings.warn("Could not extract reference from data file.") + try: + columns = ": ".join(comments[-1].split(": ")[1:]).split(", ") + xlabel = columns[0] + ylabel = columns[1] + except: + warnings.warn("Could not extract column names from data file.") + + # Prepare data and create Datablock + only_one_datapoint = len(data.shape) == 1 + if only_one_datapoint: + data = data.reshape(1, -1) + x = data[..., 0] + y = data[..., 1] + stat_unc = data[..., 2] + syst_unc = data[..., 3] + cov_stat = np.loadtxt(self.file_cov_stat, **filekwargs) + cov_syst = np.loadtxt(self.file_cov_syst, **filekwargs) + if only_one_datapoint: + cov_stat = cov_stat.reshape((1, 1)) + cov_syst = cov_syst.reshape((1, 1)) + if cov_stat.shape != (len(x), len(x)): + raise ValueError( + f"Statistical covariance matrix must be of shape {(len(x), len(x))}, but got {cov_stat.shape}" + ) + if cov_syst.shape != (len(x), len(x)): + raise ValueError( + f"Systematic covariance matrix must be of shape {(len(x), len(x))}, but got {cov_syst.shape}" + ) + if not np.allclose(stat_unc, np.sqrt(np.diag(cov_stat))): + raise ValueError( + "Statistical uncertainties (squared) do not match the diagonal of the statistical covariance matrix" + ) + if not np.allclose(syst_unc, np.sqrt(np.diag(cov_syst))): + raise ValueError( + "Systematic uncertainties (squared) do not match the diagonal of the systematic covariance matrix" + ) + return Datablock( + x, + y, + xlabel, + ylabel, + cov_stat=cov_stat, + cov_syst=cov_syst, + metadata=metadata, + ) diff --git a/src/dispersionrelations/dynamics.py b/src/dispersionrelations/dynamics.py index 3eeccdc..a302471 100644 --- a/src/dispersionrelations/dynamics.py +++ b/src/dispersionrelations/dynamics.py @@ -1,9 +1,37 @@ +from abc import ABC, abstractmethod +import warnings +from typing import Union, Optional # for python<=3.9 + import numpy as np from scipy.interpolate import interp1d -from dispersionrelations.constants import * -from dispersionrelations.utils import * -from dispersionrelations.kinematics import * -from dispersionrelations.integrals import * + +from dispersionrelations.constants import GeV, alpha_fs, E_almost_zero +from dispersionrelations.kinematics import phase_space_twobody +from dispersionrelations.integrals import DispersionIntegralRHC, integrate_gl + + +def sR(M_R, G_R): + r""" + Resonance pole location calculated from its mass and width. + + Parameters + ---------- + M_R : float + Mass of the particle. + G_R : float + Width of the particle. + + Returns + ------- + sR : complex + The complex pole location. + + Notes + ----- + The pole location is calculated using :math:`s_R(M_R, \Gamma_R) = (M_R - i \Gamma_R / 2)^2`. + """ + return (M_R - 1j * G_R / 2) ** 2 + def vertex_VPP__2(s, mP, mP_2=None): @@ -262,7 +290,53 @@ def vertex_VAP(s, mA, mP): return np.sqrt(vertex_VAP__2(s, mA, mP) + 0j) -def taming_BlattWeisskopf(s, sthr, sB, l=1, normalize_at_sB=False): +def vertex_VAP_T1__2(s, mA, mP): + r"""A vector :math:`\to` axial vector–pseudoscalar vertex squared. + + Parameters + ---------- + s : array_like + Kinematic variable (can be complex). + + mV : float + Mass of the axial vector with momentum :math:`k`. + + mP : float + Mass of the pseudoscalar with momentum :math:`p`. + + Returns + ------- + b : array_like + The same shape as input `s`. + """ + return s / 3 + + +def vertex_VAP_T2__2(s, mA, mP): + r"""A vector :math:`\to` axial vector–pseudoscalar vertex squared. + + Parameters + ---------- + s : array_like + Kinematic variable (can be complex). + + mV : float + Mass of the axial vector with momentum :math:`k`. + + mP : float + Mass of the pseudoscalar with momentum :math:`p`. + + Returns + ------- + b : array_like + The same shape as input `s`. + """ + return 1 / 6 + s * 0 + + +def taming_BlattWeisskopf( + s, sthr: float, sB: float, l: int = 1, normalize_at_sB: bool = False +): r""" Blatt–Weisskopf taming factors. @@ -328,7 +402,7 @@ def taming_BlattWeisskopf(s, sthr, sB, l=1, normalize_at_sB=False): ) -def taming_spacelike_pole(s, sB=(2 * GeV) ** 2, n=1): +def taming_spacelike_pole(s, sB: float = (2 * GeV) ** 2, n: int = 1): r""" A simple decaying function, used as a taming factor. @@ -357,7 +431,7 @@ def taming_spacelike_pole(s, sB=(2 * GeV) ** 2, n=1): return (sB / (sB + s)) ** n -def BreitWigner(s, mass, width): +def BreitWigner(s, mass: float, width: float): r""" Breit–Wigner distribution. @@ -387,7 +461,7 @@ def BreitWigner(s, mass, width): return 1 / (s - mass**2 + 1j * mass * width) -def BreitWignerED(s, mass, width_s, width_0): +def BreitWignerED(s, mass: float, width_s: callable, width_0: float): r""" Breit–Wigner distribution. @@ -424,7 +498,7 @@ def BreitWignerED(s, mass, width_s, width_0): return 1 / (s - mass**2 + 1j * mass * width_s(s, mass, width_0)) -def radiative_width_to_normalisation(width, mV, mP): +def radiative_width_to_normalisation_squared(width, mV, mP): r"""Computes the vector–pseudoscalar form factor normalisation squared from the radiative width of the vector particle. Parameters @@ -438,7 +512,7 @@ def radiative_width_to_normalisation(width, mV, mP): Returns ------- - c : float + c_squared : float Normalisation squared :math:`|f(0)|^2`. Notes @@ -452,9 +526,47 @@ def radiative_width_to_normalisation(width, mV, mP): return width / ((alpha_fs / 24) * ((mV**2 - mP**2) / mV) ** 3) -class TwoBodyChannel(DispersionIntegralRHC): +def radiative_width_to_normalisation_squared_error( + width, mV, mP, width_std, mV_std, mP_std +): + r"""Computes the error of the vector–pseudoscalar form factor normalisation squared from the radiative width of the vector particle. + + Parameters + ---------- + width, width_std : float + Radiative width of the vector particle (and the corresponding error). + mV, mV_std : float + Mass of the vector particle (and the corresponding error). + mP, mP_std : float + Mass of the pseudoscalar particle (and the corresponding error). + Returns + ------- + c_squared_std : float + Error of the normalisation squared :math:`|f(0)|^2`. + + Notes + ----- + + The radiative width :math:`\Gamma` is connected to the form factor normalisation :math:`|f(0)|` via + + .. math:: \Gamma_{V \rightarrow P\gamma} = \frac{\alpha(M_V^2 - M_P^2)^3}{24 M_V^3} |f_{VP}(0)|^2. + + """ + jacobian_width = 1 / ((alpha_fs / 24) * ((mV**2 - mP**2) / mV) ** 3) + jacobian_mV = ( + 3 * width * mV**2 * (mP**2 + mV**2) / (alpha_fs * (mV**2 - mP**2) ** 4) + ) + jacobian_mP = 6 * width * mP * mV**3 / (alpha_fs * (mV**2 - mP**2) ** 4) + return np.sqrt( + (jacobian_width * width_std) ** 2 + + (jacobian_mV * mV_std) ** 2 + + (jacobian_mP * mP_std) ** 2 + ) + + +class Channel(DispersionIntegralRHC, ABC): r""" - Dispersive representation of a two-body channel. + Dispersive representation of a channel. Parameters ---------- @@ -489,15 +601,27 @@ class TwoBodyChannel(DispersionIntegralRHC): dispersionrelations.dynamics.StableTwoBodyChannel """ + def __init__( self, - loop_integrand, - threshold, - infinity=1e6, - integration_split_points=[2, 3, 5, 8, 10, 50, 1e2, 1e3, 1e4, 1e5], - integration_order=100, - subtraction_point=0, - subtraction_constants=[0], + loop_integrand: callable, + threshold: float, + infinity: float = 1e6, + integration_split_points: list[float] = [ + 2, + 3, + 5, + 8, + 10, + 50, + 1e2, + 1e3, + 1e4, + 1e5, + ], + integration_order: int = 100, + subtraction_point: float = 0, + subtraction_constants: list[float] = [0], ): super().__init__( loop_integrand, @@ -561,8 +685,57 @@ def loop(self, s, sheet=1): result -= self.loop_discontinuity(s) return result + @abstractmethod + def phase_space_function(self, s): + return None -class StableTwoBodyChannel(TwoBodyChannel): + @abstractmethod + def vertex_function_squared(self, s): + return None + + @abstractmethod + def taming_factor_squared(self, s): + return None + + def vertex_function(self, s): + return np.sqrt(self.vertex_function_squared(s) + 0j) + + def taming_factor(self, s): + return np.sqrt(self.taming_factor_squared(s) + 0j) + + +class PrecomputedChannel(Channel): + def __init__(self, s_values, integrand_values, threshold, **kwargs): + self.last_s = s_values[-1] + self.last_integrand = integrand_values[-1] + self.interpolated_integrand = interp1d(s_values, integrand_values) + warnings.warn( + "You are loading this channel from precomputed values. Full discontinuity will be absorbed in the phase space function." + ) + super().__init__(self.extrapolated_integrand, threshold, **kwargs) + + def extrapolated_integrand(self, s): + s = np.asarray(s) + res = np.zeros_like(s) + mask1 = s < self.threshold + mask2 = s > self.last_s + mask3 = ~(mask1 | mask2) + res[mask1] = 0 + res[mask2] = self.last_integrand + res[mask3] = self.interpolated_integrand(s[mask3]) + return res + + def phase_space_function(self, s): + return self.integrand(s) + + def vertex_function_squared(self, s): + return 0 * s + 1 + + def taming_factor_squared(self, s): + return 0 * s + 1 + + +class StableTwoBodyChannel(Channel): r""" Constructs a two-body channel with stable particles with masses :code:`m1` and :code:`m2`. @@ -590,15 +763,39 @@ class StableTwoBodyChannel(TwoBodyChannel): where :math:`\rho` is the phase space, :math:`\beta` is the vertex function, and :math:`B` is the taming factor. """ - def __init__(self, m1, m2, vertex_function_squared=vertex_VPP__2, taming_factor=taming_spacelike_pole, **kwargs): + + def __init__( + self, + m1, + m2, + vertex_function_squared=vertex_VPP__2, + taming_factor_squared=lambda s: taming_spacelike_pole(s) ** 2, + **kwargs, + ): + # Store parameters: self.m1 = m1 self.m2 = m2 - self.threshold = (m1 + m2)**2 - self.phase_space = lambda s: phase_space_twobody(s, m1, m2) - self.vertex_function_squared = lambda s: vertex_function_squared(s, m1, m2) - self.taming_factor = taming_factor - self.loop_integrand = lambda s: self.phase_space(s) * self.vertex_function_squared(s) * self.taming_factor(s)**2 - super().__init__(self.loop_integrand, self.threshold, **kwargs) + self.threshold = (m1 + m2) ** 2 + # Implement abstract methods: + self.__psf__ = np.vectorize(lambda s: phase_space_twobody(s, m1, m2)) + self.__vfs__ = np.vectorize(lambda s: vertex_function_squared(s, m1, m2)) + self.__tfs__ = np.vectorize(lambda s: taming_factor_squared(s)) + # Construct the integrand and call the parent constructor: + loop_integrand = ( + lambda s: self.phase_space_function(s) + * self.vertex_function_squared(s) + * self.taming_factor_squared(s) + ) + super().__init__(loop_integrand, self.threshold, **kwargs) + + def phase_space_function(self, s): + return self.__psf__(s) + + def vertex_function_squared(self, s): + return self.__vfs__(s) + + def taming_factor_squared(self, s): + return self.__tfs__(s) class UnstableParticle: @@ -613,7 +810,7 @@ class UnstableParticle: as long as it has a :code:`threshold` and functions :code:`integrand` and :code:`loop`. pole_location : complex - Location of the resonance pole (see :func:`dispersionrelations.constants.sR`). + Location of the resonance pole (see :func:`sR`). pole_sheet : int Riemann sheet of the decay channel, on which the resonance lives. For above-threshold resonances, the typical sheet is `2`. @@ -625,12 +822,13 @@ class UnstableParticle: interpolation_density : int Number of points per interpolation interval. """ + def __init__( self, - decay_channel, - pole_location, - pole_sheet=2, - mass_and_coupling=None, + decay_channel: Channel, + pole_location: complex, + pole_sheet: int = 2, + mass_and_coupling: Optional[tuple] = None, interpolation_points=[1, 2, 5, 1e1, 5e1, 1e2, 1e3, 1e5, 1e8], interpolation_density=1000, ): @@ -797,25 +995,45 @@ class SemiStableTwoBodyCut: interpolation_density: int Number of points per interpolation interval. """ + def __init__( self, - unstable_particle, - spectator_mass, - vertex_function, - phase_space=phase_space_twobody, - taming_factor=taming_spacelike_pole, - gl_order=100, - split_n=3, - interpolation_points=[1, 2, 3, 4, 5, 1e1, 3e1, 5e1, 1e2, 1e3, 1e5, 1e8, 1e9], - interpolation_density=1000, + unstable_particle: UnstableParticle, + spectator_mass: float, + vertex_function_squared: callable, + gl_order: int = 100, + split_n: int = 3, + integration_split_points: list[float] = [2, 5, 10, 50, 100], + interpolation_points: list = [ + 1, + 2, + 3, + 4, + 5, + 1e1, + 3e1, + 5e1, + 1e2, + 1e3, + 1e4, + 1e5, + 1e6, + 1e7, + 1e8, + 1e9, + ], + interpolation_density: int = 1000, ): self.unstable_particle = unstable_particle self.spectator_mass = spectator_mass - self.phase_space = phase_space - self.vertex_function = vertex_function - self.taming_factor = taming_factor + self.phase_space_factor = phase_space_twobody + self.vertex_function_squared = vertex_function_squared self.gl_order = gl_order self.split_n = split_n + self.integration_split_points = ( + self.unstable_particle.decay_channel.threshold + * np.array(integration_split_points) + ) self.interpolation_points = interpolation_points self.interpolation_density = interpolation_density self.__interpolate_ImPI_integral__() @@ -847,9 +1065,8 @@ def ImPI_integrand_RE(self, s, x): """ return ( self.unstable_particle.spectral_function_RE(x) - * self.phase_space(s, np.sqrt(x + 0j), self.spectator_mass) - * self.vertex_function(s, np.sqrt(x + 0j), self.spectator_mass) ** 2 - * self.taming_factor(s) ** 2 + * self.phase_space_factor(s, np.sqrt(x + 0j), self.spectator_mass) + * self.vertex_function_squared(s, np.sqrt(x + 0j), self.spectator_mass) ) def ImPI_integrand_CP(self, s, x): @@ -879,9 +1096,8 @@ def ImPI_integrand_CP(self, s, x): """ return ( self.unstable_particle.spectral_function_CP(x) - * self.phase_space(s, np.sqrt(x + 0j), self.spectator_mass) - * self.vertex_function(s, np.sqrt(x + 0j), self.spectator_mass) ** 2 - * self.taming_factor(s) ** 2 + * self.phase_space_factor(s, np.sqrt(x + 0j), self.spectator_mass) + * self.vertex_function_squared(s, np.sqrt(x + 0j), self.spectator_mass) ) def ImPI_integral(self, s): @@ -918,14 +1134,33 @@ def ImPI_integral(self, s): if np.imag(upper_bound) != 0: integrand = lambda x: self.ImPI_integrand_CP(s, x) - integral_1 = integrate_gl( - integrand, - lower_bound, - middle_bound, - order=self.gl_order, - split_n=self.split_n, + integration_points_1 = np.concatenate( + ( + [lower_bound], + self.integration_split_points[ + self.integration_split_points < middle_bound + ], + [middle_bound], + ) ) + integral_1 = 0 * s + 0j + for i in range(len(integration_points_1) - 1): + integral_1 += integrate_gl( + integrand, + integration_points_1[i], + integration_points_1[i + 1], + order=self.gl_order, + split_n=self.split_n, + ) + # integral_1 = integrate_gl( + # integrand, + # lower_bound, + # middle_bound, + # order=self.gl_order, + # split_n=self.split_n, + # ) + if np.imag(upper_bound) == 0: return integral_1 @@ -995,8 +1230,8 @@ def ImPI_integral_RE(self, s): s[interpolatable] ) if np.sum(above_interpolation_maximum) > 0: - result[~above_interpolation_maximum] = self.ImPI_integral_vectorized( - s[~above_interpolation_maximum] + result[above_interpolation_maximum] = self.ImPI_integral_vectorized( + s[above_interpolation_maximum] ) return result @@ -1043,3 +1278,236 @@ def __call__(self, s): result[~is_complex] = self.ImPI_integral_RE(np.real(s[~is_complex])) return result + +class SemiStableTwoBodyChannel(Channel): + def __init__( + self, + twobodycut: SemiStableTwoBodyCut, + taming_factor_squared=lambda s: taming_spacelike_pole(s) ** 2, + ): + # Store parameters: + threshold = ( + np.sqrt(twobodycut.unstable_particle.decay_channel.threshold) + + twobodycut.spectator_mass + ) ** 2 + # Implement abstract methods: + self.__psf__ = twobodycut.__call__ + self.__vfs__ = np.vectorize(lambda s: 1) + self.__tfs__ = np.vectorize(taming_factor_squared) + # Construct the integrand and call the parent constructor: + loop_integrand = ( + lambda s: self.phase_space_function(s) + * self.vertex_function_squared(s) + * self.taming_factor_squared(s) + ) + super().__init__(loop_integrand, threshold) + + def phase_space_function(self, s): + return self.__psf__(s) + + def vertex_function_squared(self, s): + return self.__vfs__(s) + + def taming_factor_squared(self, s): + return self.__tfs__(s) + + +class TwoPotentialModel: + def __init__( + self, + channels: dict[str, Channel], + bg_par: dict, + res_par: dict, + channel_sheets: dict = {}, + s_array: Union[np.ndarray, list] = np.linspace(E_almost_zero, 3 * GeV, 200) + ** 2, + ): + # Initialize channels + self.channels = channels + self.channel_sheets = {key: 1 for key in channels.keys()} + for key in channel_sheets.keys(): + if key in channels.keys(): + self.channel_sheets[key] = channel_sheets[key] + else: + warnings.warn( + f"Channel '{key}' not found in the provided channels dictionary. Ignoring the sheet specification." + ) + # Initialize parameters + self.bg_par = bg_par + self.res_par = res_par + # Initialize auxilary variables + self.nC = len(channels) + self.nR = len(res_par["m"]) + self.IdC = np.identity(self.nC) + self.IdR = np.identity(self.nR) + + # Initialize kinematic variables + self.__set_s_array__(s_array) + + # Default parameters, in case none are provided + self.bg_par = { + "f0": np.zeros((self.nC, self.nC)), + "fR": np.array([0.0]), + } + self.res_par = { + "g": np.zeros((self.nR, self.nC)), + "m": np.arange(self.nR) + 1 * GeV, + "c": np.ones(self.nC), + "a": np.ones(self.nR), + } + + # Store model parameters from passed dictionaries + self.store_parameters(bg_par, res_par) + + # Computation steps + self.precompute() + self.compute_bg() + self.compute_res() + + def __set_s_array__(self, s_array): + # Index _rs stands for "reshaped" + self.s_array = np.array(s_array) + self.s_rs = self.s_array[..., np.newaxis, np.newaxis] + self.E_array = np.sign(s_array) * np.sqrt(np.absolute(self.s_array)) + self.E_rs = self.E_array[..., np.newaxis, np.newaxis] + + def store_parameters(self, new_bg_par=None, new_res_par=None): + if new_bg_par is not None: + for key in new_bg_par: + if key in self.bg_par.keys(): + self.bg_par[key] = new_bg_par[key] + else: + warnings.warn( + f"Background parameter '{key}' not recognized. Ignoring." + ) + + self.f0 = np.array(self.bg_par["f0"]) + self.fR = np.array(self.bg_par["fR"]) + + if new_res_par is not None: + for key in new_res_par: + if key in self.res_par.keys(): + self.res_par[key] = new_res_par[key] + else: + warnings.warn( + f"Resonance parameter '{key}' not recognized. Ignoring." + ) + + self.g = np.array(self.res_par["g"]) + self.m = np.array(self.res_par["m"]) + self.c = np.array(self.res_par["c"]) + self.a = np.array(self.res_par["a"]) + + def __diag_in_channel_space__(self, arr): + arr = np.asarray(arr) + if arr.shape != (*self.s_array.shape, self.nC): + raise ValueError( + f"Input array must have shape {(*self.s_array.shape, self.nC)}" + ) + + channel_indexing = np.arange(self.nC) + result = np.zeros((*self.s_array.shape, self.nC, self.nC)) + 0j + result[..., channel_indexing, channel_indexing] = arr + return result + + def __diag_in_resonance_space__(self, arr): + arr = np.asarray(arr) + if arr.shape != (*self.s_array.shape, self.nR): + raise ValueError( + f"Input array must have shape {(*self.s_array.shape, self.nR)}" + ) + + resonance_indexing = np.arange(self.nR) + result = np.zeros((*self.s_array.shape, self.nR, self.nR)) + 0j + result[..., resonance_indexing, resonance_indexing] = arr + return result + + def precompute(self): + loops_arr = np.array( + [ + self.channels[key].loop(self.s_array, sheet=self.channel_sheets[key]) + for key in self.channels.keys() + ] + ) + self.loops = self.__diag_in_channel_space__(np.moveaxis(loops_arr, 0, -1)) + + def compute_bg(self, new_bg_par=None): + if new_bg_par is not None: + self.store_parameters(new_bg_par=new_bg_par) + + self.f_s = ( + np.ones((*self.s_array.shape, self.nC, self.nC)) + self.fR * self.s_rs + ) + self.VB = self.f0 / self.f_s + self.TB = np.linalg.solve(self.IdC - self.VB @ self.loops, self.VB) + self.vertex = self.IdC + self.TB @ self.loops + self.loops_with_bg = self.loops @ self.vertex + + def compute_res(self, new_res_par=None): + if new_res_par is not None: + self.store_parameters(new_res_par=new_res_par) + + self.GR_inv = self.__diag_in_resonance_space__( + self.s_array[..., np.newaxis] - self.m[np.newaxis, ...] ** 2 + ) + self.GR_renormalized_inv = self.GR_inv + self.g @ self.loops_with_bg @ self.g.T + self.tR = -self.g.T @ np.linalg.solve(self.GR_renormalized_inv, self.g) # TODO: add a warning for failed inversion + self.TR = self.vertex @ self.tR @ (np.swapaxes(self.vertex, -2, -1)) + self.F = np.einsum( + "...ij,...j->...i", + -self.vertex @ self.g.T, + np.linalg.solve(self.GR_renormalized_inv, self.a), # TODO: add a warning for failed inversion + ) + + def compute_everything(self, bg_par=None, res_par=None): + self.compute_bg(bg_par) + self.compute_res(res_par) + + def new_s_array(self, s_array): + self.__set_s_array__(s_array) + self.precompute() + self.compute_bg() + self.compute_res() + + def __amplitude__(self, amp): + vertex_factors = self.__diag_in_channel_space__( + np.moveaxis( + np.array( + [ + channel.vertex_function(self.s_array) + * channel.taming_factor(self.s_array) + for channel in self.channels.values() + ] + ), + 0, + -1, + ) + ) + return vertex_factors @ amp @ vertex_factors + + def T_matrix_bg(self): + return self.__amplitude__(self.TB) + + def T_matrix_res(self): + return self.__amplitude__(self.TR) + + def T_matrix(self): + return self.__amplitude__(self.TB + self.TR) + + def S_matrix(self): + phase_space_factors = self.__diag_in_channel_space__( + np.moveaxis( + np.array( + [ + channel.phase_space_function(self.s_array) + for channel in self.channels.values() + ] + ), + 0, + -1, + ) + ) + return self.IdC + 2j * phase_space_factors * self.T_matrix() + + def form_factor(self): + return self.F \ No newline at end of file diff --git a/src/dispersionrelations/integrals.py b/src/dispersionrelations/integrals.py index ef94691..abeaedf 100644 --- a/src/dispersionrelations/integrals.py +++ b/src/dispersionrelations/integrals.py @@ -1,5 +1,7 @@ +import warnings + import numpy as np -from dispersionrelations.utils import * +from dispersionrelations.utils import logC z_gl_100, w_gl_100 = np.polynomial.legendre.leggauss(100) z_gl_500, w_gl_500 = np.polynomial.legendre.leggauss(500) @@ -14,7 +16,7 @@ } -def integrate_gl(function, a, b, order=100, split_n=3): +def integrate_gl(function : callable, a : float, b : float, order : int = 100, split_n : int = 3): r""" Gauss–Legendre integration. @@ -47,7 +49,7 @@ def integrate_gl(function, a, b, order=100, split_n=3): return np.sum(function(s_prime) * gl_weights * jacobian) -def residue(function, z0, radius=1e-2, order=100, split_n=3): +def residue(function : callable, z0 : complex, radius : float = 1e-2, order : int = 100, split_n : int = 3): r""" Residue of a complex function. @@ -98,7 +100,7 @@ class SimplePoleMonomial: NotImplementedError Any input :code:`n` other than an integer or a half-integer triggers an error. """ - def __init__(self, n): + def __init__(self, n : float): self.n = n if n % 1 == 0: if n == 0: @@ -275,7 +277,7 @@ class DispersionIntegralRHC: integration_order : int, optional Gauss–Legendre quadrature order. subtraction_point : float, optional - Point :math:`s_0` of subtraction. + Point :math:`s_0` of subtraction. Needs to be below the threshold. subtraction_constants : array_like, optional Array of :math:`f_i` subtraction constants. Defines subtraction level :math:`n`. @@ -291,13 +293,13 @@ class DispersionIntegralRHC: """ def __init__( self, - integrand, - threshold, - infinity=1e6, - integration_split_points=[2, 5, 10, 50, 100], - integration_order=100, - subtraction_point=0, - subtraction_constants=[0], + integrand : callable, + threshold : float, + infinity : float = 1e6, + integration_split_points : list[float] = [2, 5, 10, 50, 100], + integration_order : int = 100, + subtraction_point : float = 0, + subtraction_constants : list[float] = [0], ): self.integrand = integrand self.threshold = threshold @@ -309,6 +311,8 @@ def __init__( self.order = integration_order self.s0 = subtraction_point + if self.s0 >= self.threshold: + warnings.warn(f"Subtraction point s0={self.s0} needs to be below the threshold={self.threshold}.") self.f0_arr = subtraction_constants self.n_subtr = len(subtraction_constants) diff --git a/src/dispersionrelations/utils.py b/src/dispersionrelations/utils.py index ce49208..14e15b1 100644 --- a/src/dispersionrelations/utils.py +++ b/src/dispersionrelations/utils.py @@ -1,3 +1,10 @@ +import os +import re +import json +import warnings +from datetime import datetime +import warnings + import numpy as np @@ -11,12 +18,19 @@ def sqrt_custom_branch_cut(z, t_bc, sheet=1): A complex number or sequence of complex numbers. t_bc : float The angle of the branch cut in radians. + sheet : int + The sheet of the square root function. Must be 1 or 2. Returns ------- w : array_like The same shape as input `z`. - + + Raises + ------ + ValueError + If `sheet` is not 1 or 2. + Examples -------- @@ -31,16 +45,17 @@ def sqrt_custom_branch_cut(z, t_bc, sheet=1): label = r"sqrt" sqrt_plot_input = prepare_complex_plot_input(lambda z: sqrt_custom_branch_cut(z, np.pi/7), *re_im, label) _ = complex_plot_contours(sqrt_plot_input) - """ + """ + if sheet not in (1, 2): + raise ValueError("Sheet for the square root must be 1 or 2.") + sheet_selector = dict({1: 1, 2: -1})[sheet] r = np.abs(z) - if sheet==2: - r *= -1 t = np.angle(z) t = np.where(t >= t_bc, t, t + 2 * np.pi) - t_bc - return np.sqrt(r) * np.exp(1j * t / 2) + return np.sqrt(r) * np.exp(1j * t / 2) * sheet_selector -def sqrtRHC(z): +def sqrtRHC(z, sheet=1): r""" Square root with the right-hand-cut :math:`z\in[0,\infty)` in the real part. @@ -48,12 +63,19 @@ def sqrtRHC(z): ---------- z : array_like A complex number or sequence of complex numbers. + sheet : int + The sheet of the square root function. Must be 1 or 2. Returns ------- w : array_like The same shape as input `z`. + Raises + ------ + ValueError + If `sheet` is not 1 or 2. + Examples -------- @@ -69,10 +91,10 @@ def sqrtRHC(z): sqrtRHC_plot_input = prepare_complex_plot_input(sqrtRHC, *re_im, label) _ = complex_plot_contours(sqrtRHC_plot_input) """ - return sqrt_custom_branch_cut(z, t_bc=0) + return sqrt_custom_branch_cut(z, t_bc=0, sheet=sheet) -def sqrtLHC(z): +def sqrtLHC(z, sheet=1): r""" Square root with the left-hand-cut :math:`z\in(-\infty,0]` in the real part. @@ -80,12 +102,19 @@ def sqrtLHC(z): ---------- z : array_like A complex number or sequence of complex numbers. + sheet : int + The sheet of the square root function. Must be 1 or 2. Returns ------- w : array_like The same shape as input `z`. + Raises + ------ + ValueError + If `sheet` is not 1 or 2. + Examples -------- @@ -112,7 +141,7 @@ def sqrtLHC(z): sqrt_plot_input = prepare_complex_plot_input(np.sqrt, *re_im, label) _ = complex_plot_contours(sqrt_plot_input) """ - return sqrt_custom_branch_cut(z, t_bc=np.pi) + return sqrt_custom_branch_cut(z, t_bc=np.pi, sheet=sheet) def log_custom_branch_cut(z, t_bc): @@ -347,275 +376,186 @@ def conformal_variable(s, sE, sL): return (sqrt_A - sqrt_B) / (sqrt_A + sqrt_B) -def chi2(fit_data, observed_data, covariance_matrix, weights=None): +def cite(ref): r""" - :math:`\chi^2`, computed using modelled and observed data. + Produces a LaTeX citation command. Parameters ---------- - fit_data : array_like - Data obtained from a model, :math:`f_i \in \{f_1, f_2, \dots, f_n\}`. - observed_data : array_like - Data obtained from an experiment, :math:`o_i \in \{o_1, o_2, \dots, o_n\}`. - covariance_matrix : array_like - Covariance matrix of the experimental data, :math:`C \in \mathbb{R}^{n\times n}`. - weights : array_like - Weight matrix, used to manually attenuate parts of the data, :math:`W \in \mathbb{R}^{n\times n}`. + ref : str + The reference key. Returns ------- - s : float - The total :math:`\chi^2`. - - Notes - ----- - The :math:`\chi^2` is calculated using + citation : str + The LaTeX citation command :code:`\cite{ref}`. - .. math:: \chi^2 = \sum_i \sum_j (f_i-o_i)(C^{-1}_{ij}W_{ij})(f_j-o_j). - - See also + See Also -------- - dispersionrelations.utils.chi2_vector : for individual contributions. + uncite : Inverse operation. """ - if weights is not None: - covariance_matrix = covariance_matrix / weights - difference = fit_data - observed_data - return difference @ np.linalg.inv(covariance_matrix) @ difference + return r"\cite{" + ref + "}" -def chi2_with_inverse(fit_data, observed_data, inverse_covariance_matrix, weights=None): +def uncite(citation): r""" - :math:`\chi^2`, computed using modelled and observed data. - This version of the function accepts the inverse of the covariance matrix - as a parameter and therefore gains speed by avoiding matrix inversion. + Inverse operation of :func:`cite`. + Extracts the reference key from a LaTeX citation command. Parameters ---------- - fit_data : array_like - Data obtained from a model, :math:`f_i \in \{f_1, f_2, \dots, f_n\}`. - observed_data : array_like - Data obtained from an experiment, :math:`o_i \in \{o_1, o_2, \dots, o_n\}`. - inverse_covariance_matrix : array_like - Inverted covariance matrix of the experimental data, :math:`I \in \mathbb{R}^{n\times n}`. - weights : array_like - Weight matrix, used to manually attenuate parts of the data, :math:`W \in \mathbb{R}^{n\times n}`. + citation : str + The LaTeX citation command :code:`\cite{ref}`. Returns ------- - s : float - The total :math:`\chi^2`. - - Notes - ----- - The :math:`\chi^2` is calculated using + ref : str + The reference key. - .. math:: \chi^2 = \sum_i \sum_j (f_i-o_i)(I_{ij}W_{ij})(f_j-o_j). + Raises + ------ + ValueError + If the input is not in the correct format. See also -------- - dispersionrelations.utils.chi2_vector_with_inverse : for individual contributions. + cite : Forward operation. """ - if weights is not None: - inverse_covariance_matrix = inverse_covariance_matrix * weights - difference = fit_data - observed_data - return difference @ inverse_covariance_matrix @ difference + if citation[:6] != r"\cite{" or citation[-1] != "}": + raise ValueError(f"Invalid citation format: {citation}") + result = citation[6:-1] + if result == "": + warnings.warn("Empty citation key.") + return result -def chi2_vector(fit_data, observed_data, covariance_matrix, weights=None): +def save_to_json(data, filename : str = None, folder : str = "", metadata : dict = {}): r""" - :math:`\chi^2` vector. + Saves a dictionary to a JSON file. Parameters ---------- - fit_data : array_like - Data obtained from a model, :math:`f_i \in \{f_1, f_2, \dots, f_n\}`. - observed_data : array_like - Data obtained from an experiment, :math:`o_i \in \{o_1, o_2, \dots, o_n\}`. - covariance_matrix : array_like - Covariance matrix of the experimental data, :math:`C \in \mathbb{R}^{n\times n}`. - weights : array_like - Weight matrix, used to manually attenuate parts of the data, :math:`W \in \mathbb{R}^{n\times n}`. + data : dict + The data to be saved. + filename : str + The name of the file to save the data to. If None, the current date and time will be used. + folder : str + The folder to save the file in. Default is the current folder. + metadata : dict + Additional metadata to be included in the JSON file. Default is an empty dictionary. Returns ------- - v : array_like - A vector with shape (n,) of individual :math:`\chi^2` contributions. - - Notes - ----- - The :math:`\chi^2` vector is built using - - .. math:: v_i = \sum_j (f_i-o_i)(C^{-1}_{ij}W_{ij})(f_j-o_j). - - Warning - ------- - When the covariance matrix is non-diagonal, the individual components of this - vector do not have a statistical meaning, since the observations are correlated. - Use with caution! - - Examples - -------- - >>> import numpy as np - >>> from dispersionrelations.utils import chi2_vector - >>> np.random.seed(137) - >>> n = 10 - >>> x_f = np.linspace(1, 2, n) - >>> x_o = x_f + 0.1 * np.random.rand(n) - >>> x_e = 0.1 * (np.random.rand(n) + 1) / 2 - >>> x_C = np.diag(x_e)**2 - >>> print(chi2_vector(x_f, x_o, x_C)) - [0.89280398 0.01007333 0.62584069 0.67718574 0.32026806 0.99981995 0.91740446 0.01387227 2.33635045 0.18841755] - - .. plot:: + None + """ - import numpy as np - from dispersionrelations.utils import chi2_vector - import matplotlib.pyplot as plt + if filename is None: + filename = datetime.now().strftime("%Y_%m%d_%H%M%S") - np.random.seed(137) - n = 10 - x_f = np.linspace(1, 2, n) - x_o = x_f + 0.1 * np.random.rand(n) - x_e = 0.1 * (np.random.rand(n) + 1) / 2 - x_C = np.diag(x_e)**2 - x_chi2_vector = chi2_vector(x_f, x_o, x_C) - plt.errorbar(np.arange(n), x_o, yerr=x_e, fmt="o", markersize=0, capsize=2, label="observed_data", zorder=2) - plt.plot(np.arange(n), x_f, label="fit_data", zorder=2) - plt.fill_between(np.arange(n), 0.1 * x_chi2_vector, 0, color='tab:red', label=r"$\chi^2/10$", zorder=3) - plt.legend() + if os.path.exists(os.path.join(folder, filename + ".json")): + warnings.warn(f"File {filename} already exists in folder {folder}. Renaming the old file to avoid overwriting.") + os.rename(os.path.join(folder, filename + ".json"), os.path.join(folder, filename + "_old.json")) - See also - -------- - dispersionrelations.utils.chi2 - """ - if weights is not None: - covariance_matrix = covariance_matrix / weights - difference = fit_data - observed_data - return difference * (np.linalg.inv(covariance_matrix) @ difference) + with open(os.path.join(folder, filename + ".json"), "w") as f: + json.dump({**data, **metadata}, f, indent=4, ensure_ascii=False) -def chi2_vector_with_inverse( - fit_data, observed_data, inverse_covariance_matrix, weights=None -): +def load_from_json(filename : str, folder : str = ""): r""" - :math:`\chi^2`, computed using modelled and observed data. - This version of the function accepts the inverse of the covariance matrix - as a parameter and therefore gains speed by avoiding matrix inversion. + Loads a dictionary from a JSON file. Parameters ---------- - fit_data : array_like - Data obtained from a model, :math:`f_i \in \{f_1, f_2, \dots, f_n\}`. - observed_data : array_like - Data obtained from an experiment, :math:`o_i \in \{o_1, o_2, \dots, o_n\}`. - inverse_covariance_matrix : array_like - Inverted covariance matrix of the experimental data, :math:`I \in \mathbb{R}^{n\times n}`. - weights : array_like - Weight matrix, used to manually attenuate parts of the data, :math:`W \in \mathbb{R}^{n\times n}`. + filename : str + The name of the file to load the data from. + folder : str + The folder to load the file from. Default is the current folder. Returns ------- - v : array_like - A vector with shape (n,) of individual :math:`\chi^2` contributions. - - Notes - ----- - The :math:`\chi^2` vector is built using + data : dict + The loaded data. + """ + with open(os.path.join(folder, filename + ".json"), "r") as f: + return json.load(f) - .. math:: v_i = \sum_j (f_i-o_i)(I_{ij}W_{ij})(f_j-o_j). - See also - -------- - dispersionrelations.utils.chi2_with_inverse - """ - if weights is not None: - inverse_covariance_matrix = inverse_covariance_matrix * weights - difference = fit_data - observed_data - return difference * (inverse_covariance_matrix @ difference) -def chi2_reduced( - fit_data, observed_data, covariance_matrix, weights=None, number_of_parameters=0 -): +def scientific_notation(num, rounding=2): r""" - A reduced :math:`\chi^2`, computed using modelled and observed data. + Scientific notation. Parameters ---------- - fit_data : array_like - Data obtained from a model, :math:`f_i \in \{f_1, f_2, \dots, f_n\}`. - observed_data : array_like - Data obtained from an experiment, :math:`o_i \in \{o_1, o_2, \dots, o_n\}`. - covariance_matrix : array_like - Covariance matrix of the experimental data, :math:`C \in \mathbb{R}^{n\times n}`. - weights : array_like - Weight matrix, used to manually attenuate parts of the data, :math:`W \in \mathbb{R}^{n\times n}`. - number_of_parameters : int - Number of parameters in the model, used for calculating the degrees of freedom. + num : float + A real number. + rounding : int + Number of digits after the period. Returns ------- - s : float - The reduced :math:`\chi^2`. - - Notes - ----- - The reduced :math:`\chi^2` is calculated using - - .. math:: \bar{\chi}^2 = \frac{\chi^2}{\text{d.o.f.}} = \frac{1}{n - n_\text{par}} \sum_i \sum_j (f_i-o_i)(C^{-1}_{ij}W_{ij})(f_j-o_j). + num_not : str + A LaTeX code for the number representation. - See also + Examples -------- - dispersionrelations.utils.chi2 + >>> import numpy as np + >>> from DispersionRelations.constants import scientific_notation + >>> print(scientific_notation(np.pi)) + '3.1415926536' + >>> print(scientific_notation(13.4)) + '1.34 \\times 10^{1}' """ - degrees_of_freedom = len(observed_data) - number_of_parameters - return ( - chi2(fit_data, observed_data, covariance_matrix, weights) / degrees_of_freedom - ) + exponent = int(np.floor(np.log10(np.absolute(num)))) + radical = round(num / (10**exponent), rounding) + if exponent == 0: + return str(radical) + return str(radical) + r" \times 10^{" + str(exponent) + "}" -def chi2_reduced_with_inverse( - fit_data, - observed_data, - inverse_covariance_matrix, - weights=None, - number_of_parameters=0, -): +def rounding_PDG(mean, std): r""" - A reduced :math:`\chi^2`, computed using modelled and observed data. - This version of the function accepts the inverse of the covariance matrix - as a parameter and therefore gains speed by avoiding matrix inversion. + Rounding with PDG rules :cite:`Zyla:2020zbs`. Parameters ---------- - fit_data : array_like - Data obtained from a model, :math:`f_i \in \{f_1, f_2, \dots, f_n\}`. - observed_data : array_like - Data obtained from an experiment, :math:`o_i \in \{o_1, o_2, \dots, o_n\}`. - inverse_covariance_matrix : array_like - Inverted covariance matrix of the experimental data, :math:`I \in \mathbb{R}^{n\times n}`. - weights : array_like - Weight matrix, used to manually attenuate parts of the data, :math:`W \in \mathbb{R}^{n\times n}`. - number_of_parameters : int - Number of parameters in the model, used for calculating the degrees of freedom. + mean : float + Mean value of the quantity. + std : float + Standard deviation of the quantity. Returns ------- - s : float - The reduced :math:`\chi^2`. - - Notes - ----- - The reduced :math:`\chi^2` is calculated using + mean_out, std_out : (float, float) + Mean value and standard deviation, cited according to the PDG prescription. - .. math:: \bar{\chi}^2 = \frac{\chi^2}{\text{d.o.f.}} = \frac{1}{n - n_\text{par}} \sum_i \sum_j (f_i-o_i)(I_{ij}W_{ij})(f_j-o_j). - - See also + Examples -------- - dispersionrelations.utils.chi2 + >>> from DispersionRelations.constants import rounding_PDG + >>> print(rounding_PDG(0.827, 0.119)) + (0.83, 0.12) + >>> print(rounding_PDG(0.827, 0.367)) + (0.8, 0.4) """ - degrees_of_freedom = len(observed_data) - number_of_parameters - return ( - chi2_with_inverse(fit_data, observed_data, inverse_covariance_matrix, weights) - / degrees_of_freedom + before_zero = int(np.floor(np.log10(np.absolute(std)))) + 1 + three_highest_digits = int(std / 10 ** (before_zero - 3)) + to_keep = 2 # default + if 100 <= three_highest_digits <= 354: + to_keep = 2 + if 355 <= three_highest_digits <= 949: + to_keep = 1 + if 950 <= three_highest_digits <= 999: + to_keep = 1 # because 1000 has 3 digits now + three_highest_digits = 1000 + + std_out = round( + three_highest_digits * 10 ** (before_zero - 3), to_keep - before_zero ) + mean_out = round(mean, to_keep - before_zero) + + if mean_out % 1 == 0 and std_out % 1 == 0: + mean_out = int(mean_out) + std_out = int(std_out) + + return mean_out, std_out \ No newline at end of file diff --git a/tests/test_constants.py b/tests/test_constants.py index c04397d..95ae18d 100644 --- a/tests/test_constants.py +++ b/tests/test_constants.py @@ -1,7 +1,6 @@ import numpy as np from dispersionrelations.constants import * - def test_energy_unit_conversion(): assert GeV / MeV == 1e3 assert MeV / eV == 1e6 @@ -10,29 +9,3 @@ def test_energy_unit_conversion(): assert 90 * degrees == np.pi / 2 -def test_scientific_notation(): - assert scientific_notation(np.pi) == "3.14" - assert scientific_notation(np.pi, 10) == "3.1415926536" - assert scientific_notation(13.4) == r"1.34 \times 10^{1}" - assert scientific_notation(13.4, 10) == r"1.34 \times 10^{1}" - assert scientific_notation(0.134) == r"1.34 \times 10^{-1}" - assert scientific_notation(0.000134) == r"1.34 \times 10^{-4}" - assert scientific_notation(13495) == r"1.35 \times 10^{4}" - assert scientific_notation(13495, 20) == r"1.3495 \times 10^{4}" - - -def test_rounding(): - a1, b1 = rounding_PDG(0.827, 0.367) - assert a1 == 0.8 and b1 == 0.4 - a2, b2 = rounding_PDG(0.827, 0.119) - assert a2 == 0.83 and b2 == 0.12 - - -def test_sR(): - M1 = 500 - G1 = 100 - s1 = sR(M1, G1) - assert np.real(np.sqrt(s1)) == M1 - assert np.imag(np.sqrt(s1)) == -G1 / 2 - assert np.real(s1) == M1**2 - G1**2 / 4 - assert np.imag(s1) == -M1 * G1 diff --git a/tests/test_data.py b/tests/test_data.py new file mode 100644 index 0000000..1f4127f --- /dev/null +++ b/tests/test_data.py @@ -0,0 +1,2 @@ +import numpy as np +from dispersionrelations.data import * \ No newline at end of file diff --git a/tests/test_dynamics.py b/tests/test_dynamics.py index 88674df..86f0e97 100644 --- a/tests/test_dynamics.py +++ b/tests/test_dynamics.py @@ -5,4 +5,14 @@ def test_thresholds(): M2 = 0.2 STHR = (M1 + M2)**2 assert vertex_VPP(STHR, M1, M2) == 0.0 - assert vertex_VVP(STHR, M1, M2) == 0.0 \ No newline at end of file + assert vertex_VVP(STHR, M1, M2) == 0.0 + + +def test_sR(): + M1 = 500 + G1 = 100 + s1 = sR(M1, G1) + assert np.real(np.sqrt(s1)) == M1 + assert np.imag(np.sqrt(s1)) == - G1 / 2 + assert np.real(s1) == M1**2 - G1**2 / 4 + assert np.imag(s1) == - M1 * G1 \ No newline at end of file diff --git a/tests/test_utils.py b/tests/test_utils.py index 84fc63d..49b5dc7 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,14 +1,14 @@ +import pytest import numpy as np from dispersionrelations.utils import * - -def test_sqrt_along_rhc(): - sample_points = np.linspace(0, 1e4, 10) - np.testing.assert_allclose(np.sqrt(sample_points), sqrtRHC(sample_points)) - np.testing.assert_allclose(np.sqrt(sample_points), np.imag(sqrtLHC(sample_points))) +# # # # # # # # # # # # # # # # # # # # # +# HELPERS # # # # # # # # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # def assert_reflection(f, z, cut_angle=0.0, cut_in_imag=True): + """Schwarz reflection principle test for functions with branch cuts.""" z_refl = np.conj(z) rotation = np.exp(1j * cut_angle) z_rot = rotation * z @@ -21,10 +21,24 @@ def assert_reflection(f, z, cut_angle=0.0, cut_in_imag=True): np.testing.assert_allclose(f_val, -np.conj(f_val_refl)) -def test_sqrt_reflection(): - sample_points = np.linspace(-1e4, 1e4, 3) + 1j * np.linspace(-1e2, 1e2, 3) ** 2 - assert_reflection(sqrtRHC, sample_points, cut_in_imag=False) - assert_reflection(sqrtLHC, sample_points, cut_in_imag=False) +# # # # # # # # # # # # # # # # # # # # # +# TESTS # # # # # # # # # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # + +z_positive_real_values = np.linspace(0, 1e4, 24) +z_negative_real_values = np.linspace(-1e4, 0, 24) +z_real_values = np.linspace(-1e4, 1e4, 42) +z_real_values_plus_imag = z_real_values + 1j * 1e-5 +z_real_values_minus_imag = np.conj(z_real_values_plus_imag) + + +# # # # # # # # # # # # # # # # # # # # # +# sqrt_custom_branch_cut # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # + + +def test_reflection_sqrt_custom_branch_cut(): + sample_points = z_real_values_plus_imag # np.random.seed(42) test_angle = (np.random.rand() - 0.5) * 2 * np.pi assert_reflection( @@ -35,11 +49,110 @@ def test_sqrt_reflection(): ) -def test_log_reflection(): - sample_points = np.linspace(-1e4, 1e4, 4) + 1j * np.linspace(-1e2, 1e2, 4) ** 2 - assert_reflection(lambda z: logRHC(z) - 1j * np.pi, sample_points, cut_in_imag=True) - assert_reflection(lambda z: logLHC(z) - 1j * np.pi, sample_points, cut_in_imag=True) - assert_reflection(lambda z: logC(z) + 1j * np.pi, sample_points, cut_in_imag=True) +# # # # # # # # # # # # # # # # # # # # # +# sqrtRHC # # # # # # # # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # + + +def test_real_positive_values_sqrtRHC(): + sample_points = z_positive_real_values + np.testing.assert_allclose( + sqrtRHC(sample_points), + np.sqrt(sample_points), + ) + + +def test_real_positive_values_sqrtRHC_sheet_2(): + sample_points = z_positive_real_values + np.testing.assert_allclose( + sqrtRHC(sample_points, sheet=2), + -np.sqrt(sample_points), + ) + + +def test_real_negative_values_sqrtRHC(): + sample_points = z_negative_real_values + np.testing.assert_allclose( + sqrtRHC(sample_points), + 1j * np.sqrt(-sample_points + 0j), + ) + + +def test_real_negative_values_sqrtRHC_sheet_2(): + sample_points = z_negative_real_values + np.testing.assert_allclose( + sqrtRHC(sample_points, sheet=2), + -1j * np.sqrt(-sample_points + 0j), + ) + + +def test_invalid_sheets_sqrtRHC(): + with pytest.raises(ValueError): + invalid_sheets = [-1, 0, 3, None, "first", "second"] + for sheet in invalid_sheets: + _ = sqrtRHC(1.0, sheet=sheet) + + +def test_reflection_sqrtRHC(): + sample_points = z_real_values_plus_imag + assert_reflection(sqrtRHC, sample_points, cut_in_imag=False) + + +# # # # # # # # # # # # # # # # # # # # # +# sqrtLHC # # # # # # # # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # + + +def test_real_positive_values_sqrtLHC(): + sample_points = z_positive_real_values + np.testing.assert_allclose( + np.imag(sqrtLHC(sample_points)), + np.sqrt(sample_points), + ) + + +def test_real_positive_values_sqrtLHC_sheet_2(): + sample_points = z_positive_real_values + np.testing.assert_allclose( + np.imag(sqrtLHC(sample_points, sheet=2)), + -np.sqrt(sample_points), + ) + + +def test_real_negative_values_sqrtLHC(): + sample_points = z_negative_real_values + np.testing.assert_allclose( + sqrtLHC(sample_points), + -1j * np.sqrt(sample_points + 0j), + ) + + +def test_real_negative_values_sqrtLHC_sheet_2(): + sample_points = z_negative_real_values + np.testing.assert_allclose( + sqrtLHC(sample_points, sheet=2), + 1j * np.sqrt(sample_points + 0j), + ) + + +def test_invalid_sheets_sqrtLHC(): + with pytest.raises(ValueError): + invalid_sheets = [-1, 0, 3, None, "first", "second"] + for sheet in invalid_sheets: + _ = sqrtLHC(1.0, sheet=sheet) + + +def test_reflection_sqrtLHC(): + sample_points = z_real_values_plus_imag + assert_reflection(sqrtLHC, sample_points, cut_in_imag=False) + + +# # # # # # # # # # # # # # # # # # # # # +# log_custom_branch_cut # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # + +def test_reflection_log_custom_branch_cut(): + sample_points = z_real_values_plus_imag # np.random.seed(42) test_angle = (np.random.rand() - 0.5) * 2 * np.pi assert_reflection( @@ -49,8 +162,90 @@ def test_log_reflection(): cut_in_imag=True, ) +# # # # # # # # # # # # # # # # # # # # # +# logRHC # # # # # # # # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # + +def test_real_positive_values_logRHC(): + sample_points = z_positive_real_values[z_positive_real_values != 0] + np.testing.assert_allclose( + logRHC(sample_points), + np.log(sample_points), + ) + +def test_real_negative_values_logRHC(): + sample_points = z_negative_real_values[z_negative_real_values != 0] + np.testing.assert_allclose( + logRHC(sample_points), + np.log(-sample_points + 0j) + 1j * np.pi, + ) + +def test_reflection_logRHC(): + sample_points = z_real_values_plus_imag + assert_reflection(lambda z: logRHC(z) - 1j * np.pi, sample_points, cut_in_imag=True) + +# # # # # # # # # # # # # # # # # # # # # +# logLHC # # # # # # # # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # + +def test_real_positive_values_logLHC(): + sample_points = z_positive_real_values[z_positive_real_values != 0] + np.testing.assert_allclose( + logLHC(sample_points), + np.log(sample_points) + 1j * np.pi, + ) + +def test_real_negative_values_logLHC(): + sample_points = z_negative_real_values[z_negative_real_values != 0] + np.testing.assert_allclose( + logLHC(sample_points), + np.log(-sample_points + 0j), + ) + +def test_reflection_logLHC(): + sample_points = z_real_values_plus_imag + assert_reflection(lambda z: logLHC(z) - 1j * np.pi, sample_points, cut_in_imag=True) + +# # # # # # # # # # # # # # # # # # # # # +# logC # # # # # # # # # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # + +def test_real_positive_values_logC(): + sample_points = z_positive_real_values[z_positive_real_values != 0] + np.testing.assert_allclose( + logC(sample_points), + np.log(sample_points), + ) + +def test_real_negative_values_logC(): + sample_points = z_negative_real_values[z_negative_real_values != 0] + np.testing.assert_allclose( + logC(sample_points), + np.log(-sample_points + 0j) - 1j * np.pi, + ) + +def test_reflection_logC(): + sample_points = z_real_values_plus_imag + assert_reflection(lambda z: logC(z) + 1j * np.pi, sample_points, cut_in_imag=True) + + +# # # # # # # # # # # # # # # # # # # # # +# extract_phase # # # # # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # + + +def test_extract_phase_positive_start(): + E_1 = np.linspace(0, 1.1, 1000) + s_1 = E_1**2 + t_1 = 2 * np.pi * 3 * np.sin(2 * np.pi * s_1) + t_2 = 0.0 * s_1 + t_3 = -t_1 + np.testing.assert_allclose(extract_phase(np.exp(1j * t_1)), t_1) + np.testing.assert_allclose(extract_phase(np.exp(1j * t_2)), t_2) + np.testing.assert_allclose(extract_phase(np.exp(1j * t_3)), t_3) + -def test_phase_extraction(): +def test_extract_phase_negative_start(): E_1 = np.linspace(0, 1.1, 1000) s_1 = E_1**2 t_1 = 2 * np.pi * 3 * np.sin(2 * np.pi * s_1) @@ -59,3 +254,85 @@ def test_phase_extraction(): np.testing.assert_allclose(extract_phase(np.exp(1j * t_1)), t_1) np.testing.assert_allclose(extract_phase(np.exp(1j * t_2)), t_2) np.testing.assert_allclose(extract_phase(np.exp(1j * t_3)), t_3) + + +# # # # # # # # # # # # # # # # # # # # # +# cite / uncite # # # # # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # + +def test_cite_format(): + text = "reference" + cited_text = cite(text) + assert cited_text[:6] == r"\cite{" and cited_text[-1] == "}" + + +def test_cite_uncite(): + text = "reference" + cited_text = cite(text) + uncited_text = uncite(cited_text) + assert text == uncited_text + + +def test_uncite_invalid_format(): + with pytest.raises(ValueError): + invalid_citations = [ + r"text without \cite", + r"\cite{missing_end", + r"missing_start}", + r"\\cite{accidental_double_backslash}", + ] + for invalid_citation in invalid_citations: + _ = uncite(invalid_citation) + + +def test_empty_uncite(): + with pytest.warns(UserWarning): + empty_citation = r"\cite{}" + result = uncite(empty_citation) + assert result == "" + + +# # # # # # # # # # # # # # # # # # # # # +# scientific_notation # # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # + + +def test_scientific_notation(): + assert scientific_notation(np.pi) == '3.14' + assert scientific_notation(np.pi, 10) == '3.1415926536' + assert scientific_notation(13.4) == r'1.34 \times 10^{1}' + assert scientific_notation(13.4, 10) == r'1.34 \times 10^{1}' + assert scientific_notation(0.134) == r'1.34 \times 10^{-1}' + assert scientific_notation(0.000134) == r'1.34 \times 10^{-4}' + assert scientific_notation(13495) == r'1.35 \times 10^{4}' + assert scientific_notation(13495, 20) == r'1.3495 \times 10^{4}' + + +# # # # # # # # # # # # # # # # # # # # # +# PDG_rounding # # # # # # # # # # # # # +# # # # # # # # # # # # # # # # # # # # # + + +def test_PDG_rounding_case_1(): + a1, b1 = rounding_PDG(0.827, 0.367) + assert a1 == 0.8 and b1 == 0.4 + +def test_PDG_rounding_case_1_integer(): + a1, b1 = rounding_PDG(827, 367) + assert a1 == 800 and b1 == 400 + +def test_PDG_rounding_case_2(): + a2, b2 = rounding_PDG(0.827, 0.119) + assert a2 == 0.83 and b2 == 0.12 + +def test_PDG_rounding_case_2_integer(): + a2, b2 = rounding_PDG(827, 119) + assert a2 == 830 and b2 == 120 + +def test_PDG_rounding_case_3(): + a3, b3 = rounding_PDG(0.827, 0.971) + assert a3 == 0.8 and b3 == 1.0 + +def test_PDG_rounding_case_3_integer(): + a3, b3 = rounding_PDG(827, 971) + assert a3 == 800 and b3 == 1000 \ No newline at end of file