Source code for astropy.uncertainty.distributions

```# Licensed under a 3-clause BSD style license - see LICENSE.rst

"""
Built-in distribution-creation functions.
"""

from warnings import warn

import numpy as np

from astropy import units as u

from .core import Distribution

__all__ = ["normal", "poisson", "uniform"]

[docs]
def normal(
center, *, std=None, var=None, ivar=None, n_samples, cls=Distribution, **kwargs
):
"""
Create a Gaussian/normal distribution.

Parameters
----------
center : `~astropy.units.Quantity`
The center of this distribution
std : `~astropy.units.Quantity` or None
The standard deviation/σ of this distribution. Shape must match and unit
must be compatible with ``center``, or be `None` (if ``var`` or ``ivar``
are set).
var : `~astropy.units.Quantity` or None
The variance of this distribution. Shape must match and unit must be
compatible with ``center``, or be `None` (if ``std`` or ``ivar`` are set).
ivar : `~astropy.units.Quantity` or None
The inverse variance of this distribution. Shape must match and unit
must be compatible with ``center``, or be `None` (if ``std`` or ``var``
are set).
n_samples : int
The number of Monte Carlo samples to use with this distribution
cls : class
The class to use to create this distribution.  Typically a
`Distribution` subclass.

Remaining keywords are passed into the constructor of the ``cls``

Returns
-------
distr : `~astropy.uncertainty.Distribution` or object
The sampled Gaussian distribution.
The type will be the same as the parameter ``cls``.

"""
center = np.asanyarray(center)
if var is not None:
if std is None:
std = np.asanyarray(var) ** 0.5
else:
raise ValueError("normal cannot take both std and var")
if ivar is not None:
if std is None:
std = np.asanyarray(ivar) ** -0.5
else:
raise ValueError("normal cannot take both ivar and and std or var")
if std is None:
raise ValueError("normal requires one of std, var, or ivar")
else:
std = np.asanyarray(std)

randshape = np.broadcast(std, center).shape + (n_samples,)
samples = (
center[..., np.newaxis] + np.random.randn(*randshape) * std[..., np.newaxis]
)
return cls(samples, **kwargs)

COUNT_UNITS = (
u.count,
u.electron,
u.dimensionless_unscaled,
u.chan,
u.bin,
u.vox,
u.bit,
u.byte,
)

[docs]
def poisson(center, n_samples, cls=Distribution, **kwargs):
"""
Create a Poisson distribution.

Parameters
----------
center : `~astropy.units.Quantity`
The center value of this distribution (i.e., λ).
n_samples : int
The number of Monte Carlo samples to use with this distribution
cls : class
The class to use to create this distribution.  Typically a
`Distribution` subclass.

Remaining keywords are passed into the constructor of the ``cls``

Returns
-------
distr : `~astropy.uncertainty.Distribution` or object
The sampled Poisson distribution.
The type will be the same as the parameter ``cls``.
"""
# we convert to arrays because np.random.poisson has trouble with quantities
has_unit = False
if hasattr(center, "unit"):
has_unit = True
poissonarr = np.asanyarray(center.value)
else:
poissonarr = np.asanyarray(center)
randshape = poissonarr.shape + (n_samples,)

samples = np.random.poisson(poissonarr[..., np.newaxis], randshape)
if has_unit:
warn(
"units because they need the gain to be applied. It is "
"recommended you apply the gain to convert to e.g. electrons."
)
elif center.unit not in COUNT_UNITS:
warn(
f"Unit {center.unit} was provided to poisson, which is not one of"
f' {COUNT_UNITS}, and therefore suspect as a "counting" unit.  Ensure'
" you mean to use Poisson statistics."
)

# re-attach the unit
samples = samples * center.unit

return cls(samples, **kwargs)

[docs]
def uniform(
*,
lower=None,
upper=None,
center=None,
width=None,
n_samples,
cls=Distribution,
**kwargs,
):
"""
Create a Uniform distriution from the lower and upper bounds.

Note that this function requires keywords to be explicit, and requires
either ``lower``/``upper`` or ``center``/``width``.

Parameters
----------
lower : array-like
The lower edge of this distribution. If a `~astropy.units.Quantity`, the
distribution will have the same units as ``lower``.
upper : `~astropy.units.Quantity`
The upper edge of this distribution. Must match shape and if a
`~astropy.units.Quantity` must have compatible units with ``lower``.
center : array-like
The center value of the distribution. Cannot be provided at the same
time as ``lower``/``upper``.
width : array-like
The width of the distribution.  Must have the same shape and compatible
units with ``center`` (if any).
n_samples : int
The number of Monte Carlo samples to use with this distribution
cls : class
The class to use to create this distribution.  Typically a
`Distribution` subclass.

Remaining keywords are passed into the constructor of the ``cls``

Returns
-------
distr : `~astropy.uncertainty.Distribution` or object
The sampled uniform distribution.
The type will be the same as the parameter ``cls``.
"""
if center is None and width is None:
lower = np.asanyarray(lower)
upper = np.asanyarray(upper)
if lower.shape != upper.shape:
raise ValueError("lower and upper must have consistent shapes")
elif upper is None and lower is None:
center = np.asanyarray(center)
width = np.asanyarray(width)
lower = center - width / 2
upper = center + width / 2
else:
raise ValueError(
"either upper/lower or center/width must be given "
"to uniform - other combinations are not valid"
)

newshape = lower.shape + (n_samples,)
if lower.shape == tuple() and upper.shape == tuple():
width = upper - lower  # scalar
else:
width = (upper - lower)[:, np.newaxis]
lower = lower[:, np.newaxis]
samples = lower + width * np.random.uniform(size=newshape)

return cls(samples, **kwargs)

```