from copy import deepcopy
import astropy.units as u
import numpy as np
from exosim.models.signal import Signal
from exosim.output import Output
from exosim.tasks.task import Task
from exosim.utils import RunConfig
from exosim.utils.iterators import iterate_over_chunks
[docs]
class AddGainDrift(Task):
r"""
It adds a gain noise map to the array.
The gain Drift is modeled as a polynomial model for a spectral and time dependent modulation.
Notes
-----
This is a default class with standardised inputs and outputs.
The user can load this class and overwrite the "model" method
to implement a custom Task to replace this.
"""
def __init__(self):
"""
Parameters
----------
subexposures: :class:`~exosim.models.signal.Counts`
sub-exposures cached signal
parameters: dict
channel parameters dictionary
outputs: :class:`~exosim.output.output.Output` (optional)
output file
"""
self.add_task_param("subexposures", "sub-exposures cached signal")
self.add_task_param("parameters", "channel parameters dictionary")
self.add_task_param("output", "output file", None)
[docs]
def execute(self):
self.info("adding gain noise")
subexposures = self.get_task_param("subexposures")
parameters = self.get_task_param("parameters")
output = self.get_task_param("output")
self.model(subexposures, parameters, output)
[docs]
def model(
self,
subexposures: Signal,
parameters: dict,
output: Output | None = None,
) -> None:
"""
Apply a gain drift model to the subexposures.
This method models the gain drift as a polynomial trend based on time and wavelength.
Parameters
----------
subexposures : Signal
Sub-exposures cached signal.
parameters : dict
Dictionary containing the detector parameters for gain noise computation.
output : Optional[Output]
Output file to write information, defaults to None.
Returns
-------
None
This method does not return a value but updates the subexposures dataset
with applied gain noise.
"""
if output and issubclass(output.__class__, Output):
out_grp = output.create_group("gain noise")
else:
out_grp = None
time = subexposures.time * u.Unit(subexposures.time_units)
integration_times = subexposures.metadata["integration_times"]
if isinstance(integration_times, dict):
integration_times = integration_times["value"] * u.Unit(
integration_times["unit"]
)
integration_times.to(u.s)
# Calculating polynomial trend
order_t = parameters["detector"]["gain_coeff_order_t"]
coeff_t_min = parameters["detector"]["gain_coeff_t_min"]
coeff_t_max = parameters["detector"]["gain_coeff_t_max"]
coeff_t = RunConfig.random_generator.uniform(
low=coeff_t_min, high=coeff_t_max, size=order_t + 1
)
if out_grp is not None:
out_grp.write_list("gain coeff_t", coeff_t)
y_t = self._pol_t(time, coeff_t)
wl_ax = np.arange(0, subexposures.dataset.shape[2], 1)
order_w = parameters["detector"]["gain_coeff_order_w"]
coeff_w_min = parameters["detector"]["gain_coeff_w_min"]
coeff_w_max = parameters["detector"]["gain_coeff_w_max"]
coeff_w = RunConfig.random_generator.uniform(
low=coeff_w_min, high=coeff_w_max, size=order_w + 1
)
if out_grp is not None:
out_grp.write_list("gain coeff_w", coeff_w)
y_w = self._pol_w(wl_ax, coeff_w)
z = np.zeros([y_t.size, y_w.size])
for i, y in enumerate(y_t):
z[i] = y * y_w
if "gain_drift_amplitude" in parameters["detector"]:
amplitude = parameters["detector"]["gain_drift_amplitude"]
elif "gain_drift_amplitude_range_min" in parameters["detector"]:
amplitude_min = parameters["detector"]["gain_drift_amplitude_range_min"]
amplitude_max = parameters["detector"]["gain_drift_amplitude_range_max"]
amplitude = RunConfig.random_generator.uniform(amplitude_min, amplitude_max)
if out_grp is not None:
out_grp.write_list("gain amplitude random_seed", RunConfig.random_seed)
else:
self.error("missing amplitude definition for gain")
raise KeyError(
"gain drift needs either 'gain_drift_amplitude' or "
"'gain_drift_amplitude_range_min'/'gain_drift_amplitude_range_max' "
"in the detector configuration"
)
dinamic_range = z.max() - z.min()
z *= amplitude / dinamic_range
z += 1
if out_grp is not None:
out_grp.write_array("gain drift", z)
out_grp.write_array("coeff_t", coeff_t)
out_grp.write_scalar("amplitude", amplitude)
out_grp.write_array("coeff_w", coeff_w)
# Applying the multiplicative noise
for chunk in iterate_over_chunks(
subexposures.dataset, desc="applying gain noise"
):
chunk_data = deepcopy(subexposures.dataset[chunk])
subexposures.dataset[chunk] = (
chunk_data * z[chunk[0].start : chunk[0].stop, np.newaxis, :]
)
subexposures.output.flush()
@staticmethod
def _pol_t(tt: np.ndarray, coef: np.ndarray) -> np.ndarray:
"""
Polynomial trend calculation for time.
Parameters
----------
tt : np.ndarray
Time array.
coef : np.ndarray
Coefficients for the polynomial trend.
Returns
-------
np.ndarray
Calculated polynomial trend for time.
"""
x = ((tt - tt.min()) / (tt.max() - tt.min())).si.value
y = np.zeros(tt.size)
for i, c in enumerate(coef):
y += c * x**i
return y
@staticmethod
def _pol_w(wl: np.ndarray, coef: np.ndarray) -> np.ndarray:
"""
Polynomial trend calculation for wavelength.
Parameters
----------
wl : np.ndarray
spectral axis array [pixel unit].
coef : np.ndarray
Coefficients for the polynomial trend.
Returns
-------
np.ndarray
Calculated polynomial trend for wavelength.
"""
x = (wl - wl.min()) / (wl.max() - wl.min())
y = np.zeros(wl.size)
for i, c in enumerate(coef):
y += c * x**i
return y
@staticmethod
def _psd(f: np.ndarray, w0: float, f0: float) -> np.ndarray:
"""
Calculate the power spectral density.
Parameters
----------
f : np.ndarray
Frequency array.
w0 : float
Parameter for PSD calculation.
f0 : float
Parameter for PSD calculation.
Returns
-------
np.ndarray
Calculated power spectral density.
"""
psd = np.zeros_like(f)
psd[1:] = w0 * np.sqrt((f0 / f[1:]) ** 2 + 1)
return psd
@staticmethod
def _noise_generator(f: np.ndarray, psd: np.ndarray, out=None) -> np.ndarray:
"""
Generate noise based on provided frequency and power spectral density.
Parameters
----------
f : np.ndarray
Array of frequencies.
psd : np.ndarray
Power spectral density corresponding to the frequencies.
out : object, optional
Output object to write information, defaults to None.
Returns
-------
np.ndarray
Generated noise array.
"""
# Randomize amplitude
noise = RunConfig.random_generator.normal(0, psd * np.sqrt(f[-1]), psd.size)
if out is not None:
out.write_list("gain amplitude random_seed", RunConfig.random_seed)
# Randomize phase
phi = RunConfig.random_generator.normal(-0.5, 1, psd.size) * np.pi
if out is not None:
out.write_list("gain phase random_seed", RunConfig.random_seed)
noisef = noise * np.exp(1j * phi)
return np.fft.irfft(noisef) * np.sqrt(2 * noisef.size - 1)
