import astropy.units as u
import numpy as np
from numba import jit, prange
import exosim.output as output
from exosim.tasks.subexposures.compute_reading_scheme import ComputeReadingScheme
from exosim.tasks.subexposures.estimate_ch_jitter import EstimateChJitter
from exosim.tasks.task import Task
from exosim.utils.checks import check_units
from exosim.utils.iterators import searchsorted
[docs]
class PrepareInstantaneousReadOut(Task):
"""
This task prepares the instantaneous read out.
It calls :class:`~exosim.tasks.subexposures.compute_reading_scheme.ComputeReadingScheme` to compute the ramp sampling scheme,
and :class:`~exosim.tasks.subexposures.estimate_ch_jitter.EstimateChJitter` to scale the input pointing jitter to the focal planet pixel units.
The jittering is based on the focal plane oversampling factor.
Returns
--------
dict
readout_parameters dict
:class:`~astropy.units.Quantity`
sub-exposures integration times
"""
def __init__(self):
"""
Parameters
----------
main_parameters: dict
main parameters dict
focal_plane: :class:`~exosim.models.signal.CountsPerSecond`
channel focal plane
pointing_jitter: (:class:`~astropy.units.Quantity`, :class:`~astropy.units.Quantity`, :class:`~astropy.units.Quantity`)
Tuple containing the pointing jitter in the spatial and spectral direction expressed in units of deg, and jitter time expressed as sec.
parameters: dict
dictionary containing the channel parameters.
This is usually parsed from :class:`~exosim.tasks.load.load_options.LoadOptions`
output: str or :class:`~exosim.output.hdf5.hdf5.HDF5Output` or :class:`~exosim.output.hdf5.hdf5.HDF5OutputGroup`
output file
"""
self.add_task_param("parameters", "channel parameters dict")
self.add_task_param("main_parameters", "channel parameters dict")
self.add_task_param("focal_plane", "loaded focal plane")
self.add_task_param("pointing_jitter", "")
self.add_task_param("output_file", "output file")
[docs]
def execute(self):
parameters = self.get_task_param("parameters")
main_parameters = self.get_task_param("main_parameters")
pointing_jitter = self.get_task_param("pointing_jitter")
focal_plane = self.get_task_param("focal_plane")
output_file = self.get_task_param("output_file")
osf = focal_plane.metadata["oversampling"]
if pointing_jitter != (None, None, None):
estimateChJitter = EstimateChJitter()
jitter_spe, jitter_spa, y_jit, x_jit, jit_time = estimateChJitter(
pointing_jitter=pointing_jitter, parameters=parameters
)
new_freq = check_units(jit_time[1] - jit_time[0], "Hz")
readout_freq = check_units(parameters["readout"]["readout_frequency"], "Hz")
readout_oversampling = new_freq / readout_freq
else:
readout_oversampling = 1
computeReadingScheme = ComputeReadingScheme()
(
clock,
base_mask,
frame_sequence,
number_of_exposures,
) = computeReadingScheme(
parameters=parameters,
main_parameters=main_parameters,
readout_oversampling=readout_oversampling,
output_file=output_file,
)
# Number of frames contributing to each NDR
ndr_end_cumulative_sequence = np.ma.array(
frame_sequence.cumsum(),
mask=np.tile(np.logical_not(base_mask), number_of_exposures),
).compressed()
base_mask_start = base_mask
base_mask_start[0] = 1
base_mask_start[-2] = 0
ndr_start_cumulative_sequence = np.ma.array(
frame_sequence.cumsum(),
mask=np.tile(np.logical_not(base_mask_start), number_of_exposures),
).compressed()
ndr_integration_times = (
ndr_end_cumulative_sequence - ndr_start_cumulative_sequence
) * clock
focal_time_sampled = (focal_plane.time * focal_plane.time_units).to(u.s)
fp_time = searchsorted(
focal_time_sampled, ndr_start_cumulative_sequence * clock
)
self.store_dict = {
"number_of_exposures": number_of_exposures,
"frame_sequence": frame_sequence,
"ndr_end_cumulative_sequence": ndr_end_cumulative_sequence,
"ndr_start_cumulative_sequence": ndr_start_cumulative_sequence,
"ndr_integration_times": ndr_integration_times,
"simulation_clock": clock,
"fp_time": fp_time,
}
if pointing_jitter != (None, None, None):
self.debug("Pointing jitter found")
# check if the jitter is well sampled
# take min rms if min is not zero
jitter_rms_array = np.array(
[np.sqrt(np.mean(y_jit**2)), np.sqrt(np.mean(x_jit**2))]
)
try:
jitter_rms = min(jitter_rms_array[np.nonzero(jitter_rms_array)])
except ValueError:
self.debug("jitter rms is zero in both directions")
jitter_rms = 0
self.debug(f"jitter rms: {jitter_rms:.2f}")
# check if specific resolution is required else set it to 3 by default
try:
jitter_res = parameters["detector"]["jitter_rms_min_resolution"]
self.debug("found jitter rms minimum resolution")
except KeyError:
jitter_res = 3.0
# check if magnification is wanted or needed
mag = 1
# check if magnification is needed
if jitter_rms != 0 and jitter_rms < jitter_res:
self.debug("jitter rms < 3: focal plane resampling needed")
mag = np.ceil(jitter_res / jitter_rms).astype(int)
# check if magnification is suggested
try:
forced_mag = parameters["detector"]["jitter_resampler_mag"]
self.debug("found jitter resampler magnification")
except KeyError:
forced_mag = None
# if magnification is suggested choose between the suggested and the computed
if forced_mag:
if forced_mag < mag:
self.warning(
f"suggested jitter magnification ({forced_mag}) is too small: {mag} used instead"
)
else:
mag = forced_mag
# producing diagnostic info
spe_jit_ave = (
self._average_pointing(
jitter_spe.value,
ndr_start_cumulative_sequence,
ndr_end_cumulative_sequence,
)
* jitter_spe.unit
)
spa_jit_ave = (
self._average_pointing(
jitter_spa.value,
ndr_start_cumulative_sequence,
ndr_end_cumulative_sequence,
)
* jitter_spa.unit
)
x_jit_ave = (
self._average_pointing(
x_jit,
ndr_start_cumulative_sequence,
ndr_end_cumulative_sequence,
)
/ osf
* u.pix
)
y_jit_ave = (
self._average_pointing(
y_jit,
ndr_start_cumulative_sequence,
ndr_end_cumulative_sequence,
)
/ osf
* u.pix
)
jit_indexes = []
for start, stop in zip(
ndr_start_cumulative_sequence,
ndr_end_cumulative_sequence,
strict=False,
):
jit_indexes.append(np.arange(start, stop).astype(int))
self.store_dict.update(
{
"y_jit": y_jit / osf * u.pix,
"x_jit": x_jit / osf * u.pix,
"y_jit_averaged": y_jit_ave,
"x_jit_averaged": x_jit_ave,
"spe_jit_averaged": spe_jit_ave,
"spa_jit_averaged": spa_jit_ave,
"jit_indexes": {"se": jit_indexes},
"mag": mag,
"effective_osf": osf * mag,
}
)
for k, v in self.store_dict.items():
self.debug(f"{k}: {v}")
if issubclass(output_file.__class__, output.Output):
output_file.store_dictionary(
self.store_dict, "instantaneous_readout_params"
)
self.set_output([self.store_dict, ndr_integration_times])
@staticmethod
@jit(nopython=True, parallel=True)
def _average_pointing(
jitter: np.ndarray, start: np.ndarray, stop: np.ndarray
) -> np.ndarray:
jitter = jitter.astype(np.float64)
start = start.astype(np.int64)
stop = stop.astype(np.int64)
out = np.empty(start.shape[0], dtype=np.float64)
for i in prange(start.shape[0]):
s = start[i]
e = stop[i]
acc = 0.0
for j in range(s, e):
acc += jitter[j]
out[i] = acc / (e - s)
return out
