Source code for photutils.datasets.noise

# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Tools for including noise in simulated data.
"""

import numpy as np

from photutils.utils._deprecation import deprecated_positional_kwargs

__all__ = ['apply_poisson_noise', 'make_noise_image']




@deprecated_positional_kwargs(since='3.0', until='4.0')
def apply_poisson_noise(data, seed=None):
    """
    Apply Poisson noise to an array, where the value of each element in
    the input array represents the expected number of counts.

    Each pixel in the output array is generated by drawing a random
    sample from a Poisson distribution whose expectation value is given
    by the pixel value in the input array.

    Parameters
    ----------
    data : array_like
        The array on which to apply Poisson noise. Every pixel in the
        array must have a positive value (i.e., counts).

    seed : int, optional
        A seed to initialize the `numpy.random.BitGenerator`. If `None`,
        then fresh, unpredictable entropy will be pulled from the OS.

    Returns
    -------
    result : `~numpy.ndarray`
        The data array after applying Poisson noise. The output array
        will have the same dtype as the input ``data``.

    See Also
    --------
    make_noise_image

    Examples
    --------
    .. plot::
        :include-source:

        import matplotlib.pyplot as plt
        from photutils.datasets import (apply_poisson_noise,
                                        make_4gaussians_image)

        data1 = make_4gaussians_image(noise=False)
        data2 = apply_poisson_noise(data1, seed=0)

        # plot the images
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 8))
        ax1.imshow(data1, origin='lower')
        ax1.set_title('Original image')
        ax2.imshow(data2, origin='lower')
        ax2.set_title('Original image with Poisson noise applied')
    """
    data = np.asanyarray(data)
    if np.any(data < 0):
        msg = 'data must not contain any negative values'
        raise ValueError(msg)

    rng = np.random.default_rng(seed)

    return rng.poisson(data).astype(data.dtype)





@deprecated_positional_kwargs(since='3.0', until='4.0')
def make_noise_image(shape, distribution='gaussian', mean=None, stddev=None,
                     seed=None):
    r"""
    Make a noise image containing Gaussian or Poisson noise.

    This function simply takes random samples from a Gaussian
    or Poisson distribution with the given parameters. If you
    want to apply Poisson noise to existing sources, see the
    `~photutils.datasets.apply_poisson_noise` function.

    Parameters
    ----------
    shape : 2-tuple of int
        The shape of the output 2D image.

    distribution : {'gaussian', 'poisson'}
        The distribution used to generate the random noise:

        * ``'gaussian'``: Gaussian distributed noise.
        * ``'poisson'``: Poisson distributed noise.

    mean : float
        The mean of the random distribution. Required for both Gaussian
        and Poisson noise. The default is 0.

    stddev : float, optional
        The standard deviation of the Gaussian noise to add to the
        output image. Required for Gaussian noise and ignored for
        Poisson noise (the variance of the Poisson distribution is equal
        to its mean).

    seed : int, optional
        A seed to initialize the `numpy.random.BitGenerator`. If `None`,
        then fresh, unpredictable entropy will be pulled from the OS.

    Returns
    -------
    image : 2D `~numpy.ndarray`
        Image containing random noise.

    See Also
    --------
    apply_poisson_noise

    Examples
    --------
    .. plot::
        :include-source:

        import matplotlib.pyplot as plt
        from photutils.datasets import make_noise_image

        # make Gaussian and Poisson noise images
        shape = (100, 100)
        image1 = make_noise_image(shape, distribution='gaussian', mean=0.,
                                  stddev=5.)
        image2 = make_noise_image(shape, distribution='poisson', mean=5.)

        # plot the images
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 4))
        ax1.imshow(image1, origin='lower')
        ax1.set_title(r'Gaussian noise ($\mu=0$, $\sigma=5.$)')
        ax2.imshow(image2, origin='lower')
        ax2.set_title(r'Poisson noise ($\mu=5$)')
    """
    if mean is None:
        msg = "'mean' must be input"
        raise ValueError(msg)

    rng = np.random.default_rng(seed)

    if distribution == 'gaussian':
        if stddev is None:
            msg = "'stddev' must be input for Gaussian noise"
            raise ValueError(msg)
        image = rng.normal(loc=mean, scale=stddev, size=shape)
    elif distribution == 'poisson':
        image = rng.poisson(lam=mean, size=shape)
    else:
        msg = (f'Invalid distribution: {distribution}. Use either '
               "'gaussian' or 'poisson'.")
        raise ValueError(msg)

    return image