From a2e8db664b6eead83b30d4075c797ce42289217c Mon Sep 17 00:00:00 2001 From: iagolops Date: Thu, 18 Dec 2025 06:57:54 -0800 Subject: [PATCH 1/4] spec-z stretch added --- firecrown/likelihood/_base.py | 141 ++++++++++++++++++ firecrown/likelihood/_source.py | 11 ++ .../likelihood/number_counts/_systematics.py | 5 + 3 files changed, 157 insertions(+) diff --git a/firecrown/likelihood/_base.py b/firecrown/likelihood/_base.py index ce3606aeb..de6ce49b0 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 import parameters from firecrown.data_types import DataVector, TheoryVector @@ -807,6 +808,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( @@ -1005,6 +1007,145 @@ 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 using a gaussian kernel. + + :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") + elif 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 using a lorentzian kernel. + + :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") + elif 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 + + 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..786471440 100644 --- a/firecrown/likelihood/_source.py +++ b/firecrown/likelihood/_source.py @@ -10,17 +10,22 @@ 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 +39,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/_systematics.py b/firecrown/likelihood/number_counts/_systematics.py index a795eb24d..ae920797a 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.""" From 3aa2fd6f19ae6a333754a55e9d4b4c1694a7d4e0 Mon Sep 17 00:00:00 2001 From: iagolops Date: Mon, 5 Jan 2026 04:39:03 -0800 Subject: [PATCH 2/4] Fixing errors and typing --- firecrown/likelihood/_base.py | 65 ++++++++++++++++++--------------- firecrown/likelihood/_source.py | 1 + 2 files changed, 36 insertions(+), 30 deletions(-) diff --git a/firecrown/likelihood/_base.py b/firecrown/likelihood/_base.py index de6ce49b0..e0d44c5ce 100644 --- a/firecrown/likelihood/_base.py +++ b/firecrown/likelihood/_base.py @@ -1012,8 +1012,10 @@ def dndz_stretch_fog_gaussian( 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 using a gaussian kernel. + """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 @@ -1022,8 +1024,8 @@ def dndz_stretch_fog_gaussian( """ z = np.asarray(z) dndz = np.asarray(dndz) - c = 299792.458 # km/s - + c = 299792.458 # km/s + if sigma_v < 0.0: raise ValueError("Stretch Parameter (sigma_v) must be positive") elif sigma_v == 0.0: @@ -1032,26 +1034,28 @@ def dndz_stretch_fog_gaussian( 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_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.exp(-0.5 * (z_kernel / sigma_s) ** 2) kernel /= np.trapezoid(kernel, z) - dndz_new = fftconvolve(dndz, kernel, mode='same') + 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 using a lorentzian kernel. + """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 @@ -1060,36 +1064,38 @@ def dndz_stretch_fog_lorentzian( """ z = np.asarray(z) dndz = np.asarray(dndz) - c = 299792.458 # km/s - + c = 299792.458 # km/s + if sigma_v < 0.0: raise ValueError("Stretch Parameter (sigma_v) must be positive") elif 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 - - arg = - np.sqrt(2) * np.abs(z_kernel) / sigma_s - arg = np.clip(arg, -700, 0) # avoid overflow + 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 = 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 + 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 @@ -1098,7 +1104,7 @@ class SourceGalaxySpecZStretch( :ivar sigma_v: the spec-z stretch. """ - def __init__(self, sacc_tracer: str, kernel: str = 'gaussian') -> None: + 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 @@ -1111,23 +1117,21 @@ def __init__(self, sacc_tracer: str, kernel: str = 'gaussian') -> None: default_value=SOURCE_GALAXY_SYSTEMATIC_DEFAULT_SIGMA_V ) - if kernel == 'lorentzian': + 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 - ) + 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.""" @@ -1146,6 +1150,7 @@ 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 786471440..a7a4a100e 100644 --- a/firecrown/likelihood/_source.py +++ b/firecrown/likelihood/_source.py @@ -6,6 +6,7 @@ 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, From caa110c41eb152126a933aee364266410bcc716d Mon Sep 17 00:00:00 2001 From: Arthur Loureiro Date: Thu, 8 Jan 2026 16:41:26 +0100 Subject: [PATCH 3/4] tried to add the importing of the new spec functions --- firecrown/likelihood/number_counts/__init__.py | 10 ++++++++++ firecrown/likelihood/number_counts/_factories.py | 4 +++- 2 files changed, 13 insertions(+), 1 deletion(-) 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"), ] From 188463ec916b93205c790069f0a2ee946a5223ff Mon Sep 17 00:00:00 2001 From: iagolops Date: Thu, 8 Jan 2026 14:48:59 -0300 Subject: [PATCH 4/4] Fixing if and elif for pylint --- firecrown/likelihood/_base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/firecrown/likelihood/_base.py b/firecrown/likelihood/_base.py index e0d44c5ce..6edd9f3dc 100644 --- a/firecrown/likelihood/_base.py +++ b/firecrown/likelihood/_base.py @@ -1028,7 +1028,7 @@ def dndz_stretch_fog_gaussian( if sigma_v < 0.0: raise ValueError("Stretch Parameter (sigma_v) must be positive") - elif sigma_v == 0.0: + if sigma_v == 0.0: return z, dndz n = z.shape[0] @@ -1068,7 +1068,7 @@ def dndz_stretch_fog_lorentzian( if sigma_v < 0.0: raise ValueError("Stretch Parameter (sigma_v) must be positive") - elif sigma_v == 0.0: + if sigma_v == 0.0: return z, dndz n = z.shape[0]