"""Provides classes to generate noise for data cubes."""
from abc import ABCMeta, abstractmethod
import numpy as np
import astropy.units as U
from numpy.random._generator import Generator
from martini.beams import _BaseBeam
from martini.datacube import DataCube
class _BaseNoise(object):
"""
Abstract base class to inherit for classes implementing a noise model.
Classes inheriting from :class:`~martini.noise._BaseNoise` must implement a
method :meth:`~martini.noise._BaseNoise.generate` which receives one argument,
a :class:`~martini.datacube.DataCube` instance. This method
should return an :class:`~astropy.units.Quantity` array with the same shape as the
datacube array, and units of Jy (or compatible).
Parameters
----------
seed : int, optional
The integer to seed the random number generator.
See Also
--------
martini.noise.GaussianNoise
"""
__metaclass__ = ABCMeta
seed: int
rng: Generator
def __init__(self, seed: int = 0) -> None:
self.seed = seed
self.rng = np.random.default_rng(seed=seed)
return
@abstractmethod
def generate(
self, datacube: DataCube, beam: _BaseBeam
) -> U.Quantity[U.Jy * U.arcsec**-2]:
"""
Abstract method; create a cube containing noise.
Any random number generation should use the :attr:`~martini.noise._BaseNoise.rng`
generator.
Parameters
----------
datacube : ~martini.datacube.DataCube
This method will be called passing the :class:`~martini.datacube.DataCube`
instance as an argument; its attributes can thus be accessed here.
``datacube._array.shape`` is particularly relevant.
beam : ~martini.beams._BaseBeam
This method will be called passing the object derived from
:class:`~martini.beams._BaseBeam` (for example a
:class:`~martini.beams.GaussianBeam`) as an argument. Its attributes can thus
be accessed here. The beam size is needed to estimate the pre-convolution rms
required to obtain the desired post-convolution rms.
Returns
-------
~astropy.units.Quantity
:class:`~astropy.units.Quantity`, with dimensions of flux density.
Noise realization with size matching the
:attr:`~martini.datacube.DataCube._array`.
"""
pass # pragma: no cover
def reset_rng(self) -> None:
"""Reset the random number generator to its initial state."""
self.rng = np.random.default_rng(seed=self.seed)
return
[docs]
class GaussianNoise(_BaseNoise):
"""
Implementation of a simple Gaussian noise model.
Provides a :meth:`~martini.noise.GaussianNoise.generate` method producing a cube of
Gaussian noise.
Parameters
----------
rms : ~astropy.units.Quantity
:class:`~astropy.units.Quantity`, with dimensions of flux density per beam.
Desired root mean square amplitude of the noise field after convolution with the
beam.
seed : int, optional
Seed for random number generator.
"""
rms: U.Quantity[U.Jy * U.beam**-1]
def __init__(
self,
rms: U.Quantity[U.Jy * U.beam**-1] = 1.0 * U.Jy * U.beam**-1,
seed: int = 0,
) -> None:
self.target_rms = rms
super().__init__(seed=seed)
return
[docs]
def generate(
self, datacube: DataCube, beam: _BaseBeam
) -> U.Quantity[U.Jy * U.arcsec**-2]:
"""
Create a cube containing Gaussian noise.
Some numpy functions such as :func:`numpy.random.normal` strip units, so need to
handle them explicitly.
Parameters
----------
datacube : ~martini.datacube.DataCube
This method will be called passing the :class:`~martini.datacube.DataCube`
instance as an argument; its attributes can thus be accessed here.
``datacube._array.shape`` is particularly relevant.
beam : ~martini.beams._BaseBeam
This method will be called passing the object derived from
:class:`~martini.beams._BaseBeam` (for example a
:class:`~martini.beams.GaussianBeam`) as an argument. Its attributes can thus
be accessed here. The beam size is needed to estimate the pre-convolution rms
required to obtain the desired post-convolution rms.
Returns
-------
~astropy.units.Quantity
:class:`~astropy.units.Quantity`, with dimensions of flux density.
Noise realization with size matching the
:attr:`~martini.datacube.DataCube._array`.
"""
sig_maj = (beam.bmaj / 2 / np.sqrt(2 * np.log(2)) / datacube.px_size).to_value(
U.dimensionless_unscaled
)
sig_min = (beam.bmin / 2 / np.sqrt(2 * np.log(2)) / datacube.px_size).to_value(
U.dimensionless_unscaled
)
# Approximate rms reduction by convolution:
# new_rms * 2 * sqrt(pi) * beam_std ~ old_rms for a circular beam
# for an elliptical beam beam_std = sqrt(beam_std_maj * beam_std_min)
# Approximation turns out to be low by ~ 10%, correct for this:
rms = self.target_rms * 2.19568 * np.sqrt(np.pi * sig_maj * sig_min)
rms_unit = rms.unit
return (
self.rng.normal(scale=rms.to_value(rms_unit), size=datacube._array.shape)
* rms_unit
)