import numpy as np
from astropy.nddata import StdDevUncertainty
from ndcube import NDCube
from scipy.ndimage import binary_dilation
from punchbowl.data.units import split_ccd_array
from punchbowl.prefect import punch_task
from punchbowl.util import inpaint_nans
[docs]
def dn_to_photons(data_array: np.ndarray, gain_bottom: float = 4.9, gain_top: float = 4.9) -> np.ndarray:
"""Convert an input array from DN to photon count."""
gain = split_ccd_array(data_array.shape, gain_bottom, gain_top)
return data_array * gain
[docs]
def compute_noise(
data: np.ndarray,
bias_level: float = 100,
dark_level: float = 55.81,
gain_bottom: float = 4.9,
gain_top: float = 4.9,
read_noise_level: float = 17,
bitrate_signal: int = 16) -> tuple[np.ndarray, np.ndarray]:
"""
Generate noise based on an input data array, with specified noise parameters.
Parameters
----------
data
input data array (n x n)
bias_level
ccd bias level
dark_level
ccd dark level
gain_bottom
ccd gain (bottom side of CCD)
gain_top
ccd gain (top side of CCD)
read_noise_level
ccd read noise level
bitrate_signal
desired ccd data bit level
Returns
-------
data : np.ndarray
clipped data with the bias level added
noise : np.ndarray
computed noise array corresponding to input data and ccd/noise parameters
"""
data = data.copy().astype("long")
gain = split_ccd_array(data.shape, gain_bottom, gain_top)
# Photon / shot noise generation
data_photon = data * gain # DN to photoelectrons
sigma_photon = np.sqrt(data_photon) # Converting sigma of this
sigma = sigma_photon / gain # Converting back to DN
noise_photon = np.random.normal(scale=sigma)
# Add bias level and clip pixels to avoid overflow
data = np.clip(data + bias_level, 0, 2 ** bitrate_signal - 1)
# Dark noise generation
noise_level = dark_level * gain
noise_dark = np.random.poisson(lam=noise_level, size=data.shape) / gain
# Read noise generation
noise_read = np.random.normal(scale=read_noise_level, size=data.shape)
noise_read = noise_read / gain # Convert back to DN
# And then add noise terms directly
return data, noise_photon + noise_dark + noise_read
[docs]
def compute_uncertainty(data_array: np.ndarray,
dark_level: float = 55.81,
gain_bottom: float = 4.9,
gain_top: float = 4.9,
read_noise_level: float = 17,
) -> np.ndarray:
"""With an input data array compute a corresponding uncertainty array."""
# Convert this photon count to a shot noise
data = data_array.copy()
gain = split_ccd_array(data.shape, gain_bottom, gain_top)
# Photon / shot noise generation
data_photon = data * gain # DN to photoelectrons
sigma_photon = np.sqrt(data_photon) # Converting sigma of this
noise_photon = sigma_photon / gain # Converting back to DN
# Dark noise generation
noise_dark = dark_level
# Read noise generation
noise_read = read_noise_level / gain # Convert back to DN
# And then add noise terms directly
return np.sqrt(noise_photon**2 + noise_dark**2 + noise_read**2)
[docs]
def fill_saturated_pixels(data_object: NDCube,
saturated_pixels: np.ndarray,
row_threshold: int = 300) -> NDCube:
"""
Flag saturated pixels with neighborhood and flag in the uncertainty layer.
Parameters
----------
data_object : NDCube
input data cube
saturated_pixels : np.ndarray
mask where saturated pixels are True
row_threshold : int
if this many pixels are flagged as saturated in a row, the whole row is filled with neighborhood values
Returns
-------
NDCube
a cleaned NDCube with saturated pixels filled
"""
if saturated_pixels is not None:
saturated_pixels = binary_dilation(saturated_pixels, iterations=1)
row_counts = np.sum(saturated_pixels, axis=1)
saturated_pixels[row_counts > row_threshold, :] = True
data_object.uncertainty.array[saturated_pixels] = np.inf
image = data_object.data
image[saturated_pixels] = np.nan
filled_data = inpaint_nans(image, kernel_size=5)
data_object.data[...] = filled_data[...]
data_object.mask = saturated_pixels
return data_object
[docs]
@punch_task
def update_initial_uncertainty_task(data_object: NDCube,
dark_level: float = 55.81,
gain_bottom: float = 4.9,
gain_top: float = 4.9,
read_noise_level: float = 17,
bitrate_signal: int = 16,
saturated_pixels: np.ndarray | None = None) -> NDCube:
"""Prefect task to compute initial uncertainty."""
uncertainty_array = compute_uncertainty(data_object.data,
dark_level=dark_level,
gain_bottom=gain_bottom,
gain_top=gain_top,
read_noise_level=read_noise_level,
)
data_object.uncertainty = StdDevUncertainty(uncertainty_array)
fill_saturated_pixels(data_object, saturated_pixels)
data_object.meta.history.add_now("LEVEL1-initial_uncertainty", "Initial uncertainty computed with:")
data_object.meta.history.add_now("LEVEL1-initial_uncertainty", f"dark_level={dark_level}")
data_object.meta.history.add_now("LEVEL1-initial_uncertainty", f"gain_bottom={gain_bottom}")
data_object.meta.history.add_now("LEVEL1-initial_uncertainty", f"gain_top={gain_top}")
data_object.meta.history.add_now("LEVEL1-initial_uncertainty", f"read_noise_level={read_noise_level}")
data_object.meta.history.add_now("LEVEL1-initial_uncertainty", f"bitrate_signal={bitrate_signal}")
return data_object
