Source code for satlas2.utilities
"""
Implementation of various functions that ease the work, but do not belong in one of the other modules.
.. moduleauthor:: Wouter Gins <wouter.gins@kuleuven.be>
"""
from __future__ import annotations
from typing import Optional, Tuple, Union
import numpy as np
from numpy.typing import ArrayLike
from scipy.stats import chi2, norm, poisson
from .core import Model
__all__ = ['weightedAverage', 'poissonInterval', 'generateSpectrum']
[docs]def weightedAverage(x: ArrayLike,
sigma: ArrayLike,
axis: Optional[int] = None) -> Tuple[float, float]:
r"""Takes the weighted average of an array of values and the associated
errors. Calculates the scatter and statistical error, and returns
the greater of these two values.
Parameters
----------
x: ArrayLike
Array-like assortment of measured values, is transformed into a
1D-array.
sigma: ArrayLike
Array-like assortment of errors on the measured values, is transformed
into a 1D-array.
axis: Optional[int]
Axis over which the weighted average should be calculated
Returns
-------
Tuple[float, float]
Returns a tuple (weighted average, uncertainty), with the uncertainty
being the greater of the uncertainty calculated from the statistical
uncertainty and the scattering uncertainty.
Note
----
The formulas used are
.. math::
\left\langle x\right\rangle_{weighted} &= \frac{\sum_{i=1}^N \frac{x_i}
{\sigma_i^2}}
{\sum_{i=1}^N \frac{1}
{\sigma_i^2}}
\sigma_{stat}^2 &= \frac{1}{\sum_{i=1}^N \frac{1}{\sigma_i^2}}
\sigma_{scatter}^2 &= \frac{\sum_{i=1}^N \left(\frac{x_i-\left\langle
x\right\rangle_{weighted}}
{\sigma_i}\right)^2}
{\left(N-1\right)\sum_{i=1}^N \frac{1}{\sigma_i^2}}"""
x = np.array(x)
sigma = np.array(sigma)
Xstat = (1 / sigma**2).sum(axis=axis)
Xm = (x / sigma**2).sum(axis=axis) / Xstat
Xscatt = (((x - Xm) / sigma)**2).sum(axis=axis) / ((len(x) - 1) * Xstat)
Xstat = 1 / Xstat
return Xm, np.maximum.reduce([Xstat, Xscatt], axis=axis)**0.5
[docs]def poissonInterval(data: ArrayLike,
sigma: float = 1,
alpha: Optional[float] = None,
mean: bool = False) -> Tuple[float, float]:
"""Calculates the confidence interval
for the mean of a Poisson distribution.
Parameters
----------
data: ArrayLike
Samples of separate Poisson distributions.
sigma: float
The significance level given in equivalent sigma.
Defaults to 1-sigma.
alpha: Optional[float]
Significance level of interval. If given, *sigma* is ignored.
mean: bool
Set to True if the exact mean is given, by default False
Returns
-------
low, high: Tuple[float, float]
Lower and higher limits for the interval."""
if alpha is None:
a = (1 - norm.cdf(np.abs(sigma))) * 2
else:
a = alpha
if mean:
a = 1 - a
low, high = poisson.interval(a, data)
else:
low, high = (chi2.ppf(a / 2, 2 * data) / 2,
chi2.ppf(1 - a / 2, 2 * data + 2) / 2)
low = np.nan_to_num(low)
return low, high
[docs]def generateSpectrum(
models: Union[Model, list],
x: ArrayLike,
generator: Optional[callable] = np.random.default_rng().poisson
) -> ArrayLike:
"""Generates a dataset based on the models and x-values provided.
Parameters
----------
models : Union[Model, list]
A single Model or list of models. In case of a list, all models are summed together.
x : ArrayLike
The x values for which a y value has to be generated.
generator : callable, optional
A callable with one parameter that returns a random value based on this.
The default is a Poisson generator.
Returns
-------
ArrayLike
A same-sized array as x with values given by feeding the Model.f(x) value
to the generator.
"""
def evaluate(x):
try:
for model in models:
try:
f += model.f(x)
except UnboundLocalError:
f = model.f(x)
except TypeError:
f = models.f(x)
return f
y = evaluate(x)
y = generator(y)
return y