naive_stack.py

"""
A summarizer that simple makes a histogram of a point estimate
"""

from typing import Any, Generator

import numpy as np
import qp
from ceci.config import StageParameter as Param

from rail.core.data import QPHandle, TableHandle, TableLike
from rail.core.common_params import SharedParams, TOMOGRAPHY_ALL, TOMOGRAPHY_NONE
from rail.estimation.informer import PzInformer
from rail.estimation.summarizer import PZSummarizer


class NaiveStackInformer(PzInformer):
    """Placeholder Informer"""

    name = "NaiveStackInformer"
    entrypoint_function = "inform"  # the user-facing science function for this class
    interactive_function = "naive_stack_informer"
    config_options = PzInformer.config_options.copy()

    def _finalize_run(self) -> None:
        self.model = np.array([None])
        PzInformer._finalize_run(self)


class NaiveStackSummarizer(PZSummarizer):
    """Summarizer which stacks individual P(z)"""

    name = "NaiveStackSummarizer"
    entrypoint_function = "summarize"  # the user-facing science function for this class
    interactive_function = "naive_stack_summarizer"
    config_options = PZSummarizer.config_options.copy()
    config_options.update(
        zmin=SharedParams.copy_param("zmin"),
        zmax=SharedParams.copy_param("zmax"),
        nzbins=SharedParams.copy_param("nzbins"),
        seed=Param(int, 87, msg="random seed"),
        n_samples=Param(int, 1000, msg="Number of sample distributions to create"),
    )
    inputs = [("input", QPHandle)]
    outputs = [("output", QPHandle), ("single_NZ", QPHandle)]

    def __init__(self, args: Any, **kwargs: Any) -> None:
        super().__init__(args, **kwargs)
        self.zgrid: np.ndarray | None = None

    def summarize(
        self, input_data: qp.Ensemble, **kwargs
    ) -> QPHandle | dict[str, QPHandle]:
        """Summarizer for NaiveStack which returns multiple items

        Parameters
        ----------
        input_data : qp.Ensemble
            Per-galaxy p(z), and any ancillary data associated with it

        Returns
        -------
        QPHandle | dict[str, QPHandle]
            Ensemble with n(z), and any ancillary data
            Return type depends on `output_mode`
        """
        self.set_data("input", input_data)
        self.run()
        self.finalize()
        if len(self.outputs) == 1 or self.config.output_mode != "return":
            results = self.get_handle("output")
        # if there is more than one output and output_mode = return, return them all as a dictionary
        elif len(self.outputs) > 1 and self.config.output_mode == "return":
            results = {}
            for output in self.outputs:
                results[output[0]] = self.get_handle(output[0])
        return results

    def _setup_iterator(self) -> Generator:
        itr = self.input_iterator("input")
        for s, e, d in itr:
            yield s, e, d, np.ones(e - s, dtype=bool)

    def run(self) -> None:
        handle = self.get_handle("input", allow_missing=True)
        self._input_length = handle.size()
        iterator = self._setup_iterator()
        self.zgrid = np.linspace(
            self.config.zmin, self.config.zmax, self.config.nzbins + 1
        )
        assert self.zgrid is not None
        # Initializing the stacking pdf's
        yvals = np.zeros((1, len(self.zgrid)))
        bvals = np.zeros((self.config.n_samples, len(self.zgrid)))
        bootstrap_matrix = self._broadcast_bootstrap_matrix()

        first = True
        for s, e, test_data, mask in iterator:
            print(f"Process {self.rank} running estimator on chunk {s:,} - {e:,}")
            self._process_chunk(
                s, e, test_data, mask, first, bootstrap_matrix, yvals, bvals
            )
            first = False
        if self.comm is not None:  # pragma: no cover
            bvals, yvals = self._join_histograms(bvals, yvals)

        if self.rank == 0:
            sample_ens = qp.Ensemble(
                qp.interp, data=dict(xvals=self.zgrid, yvals=bvals)
            )
            qp_d = qp.Ensemble(qp.interp, data=dict(xvals=self.zgrid, yvals=yvals))
            self.add_data("output", sample_ens)
            self.add_data("single_NZ", qp_d)

    def _process_chunk(
        self,
        start: int,
        end: int,
        data: qp.Ensemble,
        mask: np.ndarray,
        _first: bool,
        bootstrap_matrix: np.ndarray,
        yvals: np.ndarray,
        bvals: np.ndarray,
    ) -> None:
        assert self.zgrid is not None
        pdf_vals = data.pdf(self.zgrid)
        squeeze_mask = np.squeeze(mask)
        yvals += np.expand_dims(
            np.sum(
                np.where(
                    np.isfinite(pdf_vals[squeeze_mask, :]), pdf_vals[squeeze_mask], 0.0
                ),
                axis=0,
            ),
            0,
        )
        # qp_d is the normalized probability of the stack, we need to know how many galaxies were
        for i in range(self.config.n_samples):
            bootstrap_draws = bootstrap_matrix[:, i]
            # Neither all of the bootstrap_draws are in this chunk nor the index starts at "start"
            chunk_mask = (bootstrap_draws >= start) & (bootstrap_draws < end)
            bootstrap_draws = bootstrap_draws[chunk_mask] - start
            zarr = np.where(squeeze_mask, pdf_vals.T, 0.0).T[bootstrap_draws]
            bvals[i] += np.sum(zarr, axis=0)


class NaiveStackMaskedSummarizer(NaiveStackSummarizer):
    name = "NaiveStackMaskedSummarizer"
    entrypoint_function = "summarize"  # the user-facing science function for this class
    interactive_function = "naive_stack_masked_summarizer"
    config_options = NaiveStackSummarizer.config_options.copy()
    config_options.update(
        selected_bin=Param(int, TOMOGRAPHY_NONE, msg=f"bin to use, or {TOMOGRAPHY_ALL} for all bins >=0 or {TOMOGRAPHY_NONE} for no masking"),
    )
    inputs = [("input", QPHandle), ("tomography_bins", TableHandle)]
    outputs = [("output", QPHandle), ("single_NZ", QPHandle)]

    def _setup_iterator(self) -> Generator:
        selected_bin = self.config.selected_bin
        if self.config.tomography_bins in ["none", None]:
            selected_bin = TOMOGRAPHY_NONE

        if selected_bin == TOMOGRAPHY_NONE:
            itrs = [self.input_iterator("input")]
        else:
            itrs = [
                self.input_iterator("input"),
                self.input_iterator("tomography_bins"),
            ]

        for it in zip(*itrs):
            first = True
            mask = None
            for s, e, d in it:
                if first:
                    start = s
                    end = e
                    pz_data = d
                    first = False
                else:
                    if selected_bin == TOMOGRAPHY_ALL:
                        mask = d["class_id"] >= 0
                    else:
                        mask = d["class_id"] == selected_bin
            if mask is None:
                mask = np.ones(
                    pz_data.npdf,  # pylint: disable=possibly-used-before-assignment
                    dtype=bool,
                )
            yield start, end, pz_data, mask  # pylint: disable=possibly-used-before-assignment

    def summarize(
        self, input_data: qp.Ensemble, tomo_bins: TableLike | None = None, **kwargs
    ) -> QPHandle:
        """Override the Summarizer.summarize() method to take tomo bins
        as an additional input

        Parameters
        ----------
        input_data : qp.Ensemble
            Per-galaxy p(z), and any ancillary data associated with it

        tomo_bins : TableLike | None, optional
            Tomographic bins file, by default None

        Returns
        -------
        QPHandle
            Ensemble with n(z), and any ancillary data
        """
        self.set_data("input", input_data)
        if tomo_bins is None:
            self.config.tomography_bins = None
        else:
            self.set_data("tomography_bins", tomo_bins)
        self.run()
        self.finalize()
        return self.get_handle("output")