Create PhotoZUncertainty Factory for Photo-z Shift and Stretch Implementation

[paulrogozenski]

Problem description It is useful to have an implementation that can vary the width of redshift distributions. As the stretching of the redshift distribution depends on the location of the mean of the distribution, it makes sense to have a PhotoZUncertainty kind of factory which can handle both uncertainties in the width and the mean of tomographic bins (i.e. stretching and shifting distributions). The base implementation would be in the source.py file and extend to the tracers in the weak_lensing.py and number_counts.py files.

Proposed solution In source.py:

SOURCE_GALAXY_SYSTEMATIC_DEFAULT_DELTA_Z = 0.0
SOURCE_GALAXY_SYSTEMATIC_DEFAULT_SIGMA_Z = 1.0

class SourceGalaxyPhotoZUncertainty(
    SourceGalaxySystematic[_SourceGalaxyArgsT], Generic[_SourceGalaxyArgsT]
):
    """A photo-z shift and stretch bias.

    This systematic shifts the photo-z distribution by some amount `delta_z`
    and the width of the photo-z distribution by `sigma_z` in the following form:

    `n(z) \rightarrow \frac{1}{\sigma_z} n(\frac{z - \Delta z - z_{mean}}{\sigma_z} + z_{mean})`

    The following parameters are special Updatable parameters, which means that
    they can be updated by the sampler, sacc_tracer is going to be used as a
    prefix for the parameters:

    :ivar delta_z: the photo-z shift.
    :ivar sigma_z: the photo-z stretch.
    """
    def __init__(self, sacc_tracer: str) -> None:
        """Create a PhotoZUncertainty object, using the specified tracer name.

        :param sacc_tracer: the name of the tracer in the SACC file. This is used
            as a prefix for its parameters.
        """
        super().__init__(parameter_prefix=sacc_tracer)

        self.delta_z = parameters.register_new_updatable_parameter(
            default_value=SOURCE_GALAXY_SYSTEMATIC_DEFAULT_DELTA_Z
        )
        self.sigma_z = parameters.register_new_updatable_parameter(
            default_value=SOURCE_GALAXY_SYSTEMATIC_DEFAULT_SIGMA_Z
        )


    def apply(
        self, tools: ModelingTools, tracer_arg: _SourceGalaxyArgsT
    ) -> _SourceGalaxyArgsT:
        """Apply a shift and stretch to the photo-z distribution of a source.

        :param tools: the modeling tools use to update the tracer arg
        :param tracer_arg: the original source galaxy tracer arg to which we
            apply the systematic.
        :return: a new source galaxy tracer arg with the systematic applied
        """
        dndz_interp = Akima1DInterpolator(tracer_arg.z, tracer_arg.dndz)
        # Compute mean redshift
        zmean = np.average(tracer_arg.z, weights = tracer_arg.dndz)
        if self.sigma_z <= 0.0:
              raise ValueError("Stretch Parameter must be positive")
        dndz = dndz_interp((tracer_arg.z - self.delta_z - zmean) / self.sigma_z + zmean, extrapolate=False)
        dndz /= self.sigma_z
        dndz[np.isnan(dndz)] = 0.0

        return replace(
            tracer_arg,
            dndz=dndz,
        )