Source code for rubix.core.noise

import jax.numpy as jnp
from typing import Callable, Dict
from rubix.telescope.noise.noise import (
    calculate_noise_cube,
    SUPPORTED_NOISE_DISTRIBUTIONS,
)

from rubix.logger import get_logger

from jaxtyping import Array, Float, jaxtyped
from beartype import beartype as typechecker




@jaxtyped(typechecker=typechecker)
def get_apply_noise(config: dict) -> Callable:
    """
    Get the function to apply noise to the datacube based on the configuration.

    Args:
        config (dict): Configuration dictionary.

    Returns:
        The function to apply noise to the datacube.

    Example
    -------
    >>> config = {
    ...     ...
    ...     "telescope": {
    ...         "name": "MUSE",
    ...         "psf": {"name": "gaussian", "size": 5, "sigma": 0.6},
    ...         "lsf": {"sigma": 0.5},
    ...         "noise": {"signal_to_noise": 1,"noise_distribution": "normal"},
    ...    },
    ...     ...
    ... }

    >>> from rubix.core.noise import get_apply_noise
    >>> apply_noise = get_apply_noise(config)
    >>> rubixdata = apply_noise(rubixdata)
    """
    if "noise" not in config["telescope"]:
        raise ValueError("Noise information not provided in telescope config")

    if "signal_to_noise" not in config["telescope"]["noise"]:
        raise ValueError("Signal to noise information not provided in noise config")

    if "noise_distribution" not in config["telescope"]["noise"]:
        raise ValueError(
            f"Noise distribution not provided in noise config. Currently supported distributions are: {SUPPORTED_NOISE_DISTRIBUTIONS}"
        )

    # Get the signal to noise ratio
    signal_to_noise = config["telescope"]["noise"]["signal_to_noise"]

    # Get the noise distribution
    noise_distribution = config["telescope"]["noise"]["noise_distribution"]

    logger = get_logger()

    def apply_noise(rubixdata: object) -> object:
        logger.info(
            f"Applying noise to datacube with signal to noise ratio: {signal_to_noise} and noise distribution: {noise_distribution}"
        )
        datacube = rubixdata.stars.datacube
        # Define S2n for each spaxel
        S2N = jnp.ones(datacube.shape[:2]) * signal_to_noise

        # Calculate the noise cube
        noise_cube = calculate_noise_cube(
            datacube, S2N, noise_distribution=noise_distribution
        )

        # Add noise to the datacube
        rubixdata.stars.datacube += noise_cube
        return rubixdata

    return apply_noise