diff --git a/firecrown/likelihood/_base.py b/firecrown/likelihood/_base.py index 439709193..31b730cf9 100644 --- a/firecrown/likelihood/_base.py +++ b/firecrown/likelihood/_base.py @@ -32,6 +32,7 @@ import sacc from pydantic import BaseModel, ConfigDict, Field from scipy.interpolate import Akima1DInterpolator +from scipy.signal import fftconvolve from firecrown.data_types import DataVector, TheoryVector from firecrown.modeling_tools import ModelingTools @@ -808,6 +809,7 @@ def apply( SOURCE_GALAXY_SYSTEMATIC_DEFAULT_DELTA_Z = 0.0 SOURCE_GALAXY_SYSTEMATIC_DEFAULT_SIGMA_Z = 1.0 +SOURCE_GALAXY_SYSTEMATIC_DEFAULT_SIGMA_V = 0.0 def dndz_shift_and_stretch_active( @@ -1006,6 +1008,150 @@ def create_global(self) -> PhotoZShiftandStretch: raise ValueError("PhotoZShiftandStretch cannot be global.") +def dndz_stretch_fog_gaussian( + z, + dndz, + sigma_v: float, +) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]: + """Apply the Fingers-of-God (FoG) effect on the spec-z distribution. + + This function applies a gaussian kernel to model the FoG effect + due to velocity dispersion within a redshift bin. + + :param z: the redshifts + :param dndz: the dndz + :param sigma_v: the velocity dispersion within the redshift bin in km/s + :return: the shifted and stretched dndz + """ + z = np.asarray(z) + dndz = np.asarray(dndz) + c = 299792.458 # km/s + + if sigma_v < 0.0: + raise ValueError("Stretch Parameter (sigma_v) must be positive") + if sigma_v == 0.0: + return z, dndz + + n = z.shape[0] + dz = z[1] - z[0] + + z_mean = np.trapezoid(dndz * z, z) + sigma_s = (1 + z_mean) * sigma_v / c + + z_kernel = (np.arange(n) - n // 2) * dz + kernel = np.exp(-0.5 * (z_kernel / sigma_s) ** 2) + kernel /= np.trapezoid(kernel, z) + + dndz_new = fftconvolve(dndz, kernel, mode="same") + dndz_new /= np.trapezoid(dndz_new, z) + + return z, dndz_new + + +def dndz_stretch_fog_lorentzian( + z, + dndz, + sigma_v: float, +) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]: + """Apply the Fingers-of-God (FoG) effect on the spec-z distribution. + + This function applies a lorentzian kernel to model the FoG effect + due to velocity dispersion within a redshift bin. + + :param z: the redshifts + :param dndz: the dndz + :param sigma_v: the velocity dispersion within the redshift bin in km/s + :return: the shifted and stretched dndz + """ + z = np.asarray(z) + dndz = np.asarray(dndz) + c = 299792.458 # km/s + + if sigma_v < 0.0: + raise ValueError("Stretch Parameter (sigma_v) must be positive") + if sigma_v == 0.0: + return z, dndz + + n = z.shape[0] + dz = z[1] - z[0] + + z_mean = np.trapezoid(dndz * z, z) + sigma_s = (1 + z_mean) * sigma_v / c + + z_kernel = (np.arange(n) - n // 2) * dz + arg = -np.sqrt(2) * np.abs(z_kernel) / sigma_s + arg = np.clip(arg, -700, 0) # avoid overflow + kernel = np.exp(arg) + kernel /= np.trapezoid(kernel, z) + + dndz_new = fftconvolve(dndz, kernel, mode="same") + dndz_new /= np.trapezoid(dndz_new, z) + + return z, dndz_new + + +class SourceGalaxySpecZStretch( + SourceGalaxySystematic[_SourceGalaxyArgsT], Generic[_SourceGalaxyArgsT] +): + """A spec-z convolution for the Fingers-of-God (FoG) effect. + + This systematic convolves the spec-z distribution with a kernel to model + the FoG effect. + + 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 sigma_v: the spec-z stretch. + """ + + def __init__(self, sacc_tracer: str, kernel: str = "gaussian") -> None: + """Create a SpecZShift 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. + :param kernel: which kernel to use when convolving the distribution. + """ + super().__init__(sacc_tracer) + + self.sigma_v = parameters.register_new_updatable_parameter( + default_value=SOURCE_GALAXY_SYSTEMATIC_DEFAULT_SIGMA_V + ) + + if kernel == "lorentzian": + self._transform = dndz_stretch_fog_lorentzian + else: + self._transform = dndz_stretch_fog_gaussian + + def apply(self, _: ModelingTools, tracer_arg: _SourceGalaxyArgsT): + """Apply a convolution to the spec-z distribution of a source.""" + new_z, new_dndz = self._transform(tracer_arg.z, tracer_arg.dndz, self.sigma_v) + return replace(tracer_arg, z=new_z, dndz=new_dndz) + + +class SpecZStretch(SourceGalaxySpecZStretch): + """Spec-z stretch systematic.""" + + +class SpecZStretchFactory(BaseModel): + """Factory class for SpecZStretch objects.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + type: Annotated[ + Literal["SpecZStretchFactory"], + Field(description="The type of the systematic."), + ] = "SpecZStretchFactory" + + def create(self, bin_name: str) -> SpecZStretch: + """Create a SpecZStretch object with the given tracer name.""" + return SpecZStretch(bin_name) + + def create_global(self) -> SpecZStretch: + """Create a SpecZStretch object with the given tracer name.""" + raise ValueError("SpecZStretch cannot be global.") + + class SourceGalaxySelectField( SourceGalaxySystematic[_SourceGalaxyArgsT], Generic[_SourceGalaxyArgsT] ): diff --git a/firecrown/likelihood/_source.py b/firecrown/likelihood/_source.py index 4ec27ba36..a7a4a100e 100644 --- a/firecrown/likelihood/_source.py +++ b/firecrown/likelihood/_source.py @@ -6,21 +6,27 @@ from firecrown.likelihood._base import ( SOURCE_GALAXY_SYSTEMATIC_DEFAULT_DELTA_Z, SOURCE_GALAXY_SYSTEMATIC_DEFAULT_SIGMA_Z, + SOURCE_GALAXY_SYSTEMATIC_DEFAULT_SIGMA_V, PhotoZShift, PhotoZShiftFactory, PhotoZShiftandStretch, PhotoZShiftandStretchFactory, + SpecZStretch, + SpecZStretchFactory, Source, SourceGalaxy, SourceGalaxyArgs, SourceGalaxyPhotoZShift, SourceGalaxyPhotoZShiftandStretch, + SourceGalaxySpecZStretch, SourceGalaxySelectField, SourceGalaxySystematic, SourceSystematic, Tracer, dndz_shift_and_stretch_active, dndz_shift_and_stretch_passive, + dndz_stretch_fog_gaussian, + dndz_stretch_fog_lorentzian, ) # All classes have been moved to _base.py @@ -34,14 +40,20 @@ "SourceGalaxySystematic", "SourceGalaxyPhotoZShift", "SourceGalaxyPhotoZShiftandStretch", + "SourceGalaxySpecZStretch", "SourceGalaxySelectField", "SourceGalaxy", "PhotoZShift", "PhotoZShiftFactory", "PhotoZShiftandStretch", "PhotoZShiftandStretchFactory", + "SpecZStretch", + "SpecZStretchFactory", "dndz_shift_and_stretch_active", "dndz_shift_and_stretch_passive", + "dndz_stretch_fog_gaussian", + "dndz_stretch_fog_lorentzian", "SOURCE_GALAXY_SYSTEMATIC_DEFAULT_DELTA_Z", "SOURCE_GALAXY_SYSTEMATIC_DEFAULT_SIGMA_Z", + "SOURCE_GALAXY_SYSTEMATIC_DEFAULT_SIGMA_V", ] diff --git a/firecrown/likelihood/number_counts/__init__.py b/firecrown/likelihood/number_counts/__init__.py index 3018a3de5..41b3b15a0 100644 --- a/firecrown/likelihood/number_counts/__init__.py +++ b/firecrown/likelihood/number_counts/__init__.py @@ -9,6 +9,13 @@ MagnificationBiasSystematic, PhotoZShift, PTNonLinearBiasSystematic, + PhotoZShiftandStretch, +) + +# Re-export shared factories from base module +from firecrown.likelihood._base import ( + SpecZStretch, + SpecZStretchFactory, ) __all__ = [ @@ -18,4 +25,7 @@ "ConstantMagnificationBiasSystematic", "PTNonLinearBiasSystematic", "MagnificationBiasSystematic", + "PhotoZShiftandStretch", + "SpecZStretch", + "SpecZStretchFactory", ] diff --git a/firecrown/likelihood/number_counts/_factories.py b/firecrown/likelihood/number_counts/_factories.py index 8492b5415..c8c66b72c 100644 --- a/firecrown/likelihood/number_counts/_factories.py +++ b/firecrown/likelihood/number_counts/_factories.py @@ -14,6 +14,7 @@ MagnificationBiasSystematic, PTNonLinearBiasSystematic, ) +from firecrown.likelihood._source import SpecZStretchFactory from firecrown.likelihood.weak_lensing import ( PhotoZShiftandStretchFactory, PhotoZShiftFactory, @@ -123,7 +124,8 @@ def create_global(self) -> ConstantMagnificationBiasSystematic: | LinearBiasSystematicFactory | PTNonLinearBiasSystematicFactory | MagnificationBiasSystematicFactory - | ConstantMagnificationBiasSystematicFactory, + | ConstantMagnificationBiasSystematicFactory + | SpecZStretchFactory, Field(discriminator="type", union_mode="left_to_right"), ] diff --git a/firecrown/likelihood/number_counts/_systematics.py b/firecrown/likelihood/number_counts/_systematics.py index 5264ae7e9..3fc177210 100644 --- a/firecrown/likelihood/number_counts/_systematics.py +++ b/firecrown/likelihood/number_counts/_systematics.py @@ -13,6 +13,7 @@ from firecrown.likelihood._base import ( SourceGalaxyPhotoZShift, SourceGalaxyPhotoZShiftandStretch, + SourceGalaxySpecZStretch, SourceGalaxySelectField, SourceGalaxySystematic, ) @@ -49,6 +50,10 @@ class PhotoZShiftandStretch(SourceGalaxyPhotoZShiftandStretch[NumberCountsArgs]) """Photo-z shift systematic.""" +class SpecZStretch(SourceGalaxySpecZStretch[NumberCountsArgs]): + """Spec-z stretch systematic.""" + + class SelectField(SourceGalaxySelectField[NumberCountsArgs]): """Systematic to select 3D field."""