Source code for exosim.tasks.subexposures.add_foregrounds

from copy import deepcopy

import astropy.units as u
import numpy as np

from exosim.tasks.task import Task
from exosim.utils.iterators import iterate_over_chunks, searchsorted




class AddForegrounds(Task):
    """
    It adds the foregrounds the sub exposures.

    Returns
    --------
    :class:`~exosim.models.signal.Counts`
        sub-exposure cached signal class
    """

    def __init__(self):
        """
        Parameters
        ----------
        subexposures: :class:`~exosim.models.signal.Counts`
            sub-exposures cached signal class
        frg_focal_plane: :class:`~exosim.models.signal.Signal`
            channel foreground focal plane
        integration_time: :class:`~astropy.units.Quantity`
            sub-exposures integration times

        """

        self.add_task_param("subexposures", "sub-exposures cached signal class")
        self.add_task_param(
            "frg_focal_plane",
            "channel foreground focal plane",
        )
        self.add_task_param(
            "integration_time",
            "sub-exposures integration times",
        )



    def execute(self):
        self.info("adding foregrounds")
        subexposures = self.get_task_param("subexposures")
        frg_focal_plane = self.get_task_param("frg_focal_plane")
        integration_time = self.get_task_param("integration_time")

        osf = frg_focal_plane.metadata["oversampling"]

        ndrs_time = (subexposures.time * subexposures.time_units).to(u.hr)
        frg = deepcopy(frg_focal_plane.data.astype(np.float64))
        frg_time = (frg_focal_plane.time * frg_focal_plane.time_units).to(u.hr)
        index = searchsorted(frg_time, ndrs_time)
        self.debug(f"time indexes: {index}")

        # decimate the frg to the subexposure size (take every osf-th pixel)
        frg = frg[:, ::osf, ::osf]

        for chunk in iterate_over_chunks(
            subexposures.dataset, desc="adding foregrounds"
        ):
            subexposures.dataset[chunk] = self.add_frg(
                subexposures.dataset[chunk],
                frg,
                index[chunk[0]],
                integration_time[chunk[0]].to(u.s).value,
            )
            subexposures.output.flush()
        subexposures.metadata["frg_focal_plane_time_indexes"] = index
        self.set_output(subexposures)


    @staticmethod


    def add_frg(ndrs, frg, index, integration_time):
        for t in range(ndrs.shape[0]):
            ndrs[t] = ndrs[t] + frg[index[t]] * integration_time[t]
        return ndrs