widths in shidts in WL and NC

LSSTDESC · Jul 19, 2024 · d7363cc · d7363cc
1 parent 13aabac
commit d7363cc
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 3 deletions.
diff --git a/firecrown/likelihood/number_counts.py b/firecrown/likelihood/number_counts.py
@@ -20,9 +20,11 @@
     SourceGalaxy,
     SourceGalaxyArgs,
     SourceGalaxyPhotoZShift,
+    SourceGalaxyPhotoZShiftandWidth,
     SourceGalaxySelectField,
     SourceGalaxySystematic,
     PhotoZShiftFactory,
+    PhotoZShiftandWidthFactory,
     Tracer,
 )
 from firecrown.metadata.two_point import InferredGalaxyZDist
@@ -60,6 +62,10 @@ class PhotoZShift(SourceGalaxyPhotoZShift[NumberCountsArgs]):
     """Photo-z shift systematic."""
 
 
+class PhotoZShiftandWidth(SourceGalaxyPhotoZShiftandWidth[NumberCountsArgs]):
+    """Photo-z shift systematic."""
+
+
 class SelectField(SourceGalaxySelectField[NumberCountsArgs]):
     """Systematic to select 3D field."""
 
@@ -562,6 +568,7 @@ def create_global(self) -> ConstantMagnificationBiasSystematic:
 
 NumberCountsSystematicFactory = Annotated[
     PhotoZShiftFactory
+    | PhotoZShiftandWidthFactory
     | LinearBiasSystematicFactory
     | PTNonLinearBiasSystematicFactory
     | MagnificationBiasSystematicFactory

diff --git a/firecrown/likelihood/weak_lensing.py b/firecrown/likelihood/weak_lensing.py
@@ -21,9 +21,11 @@
     SourceGalaxy,
     SourceGalaxyArgs,
     SourceGalaxyPhotoZShift,
+    SourceGalaxyPhotoZShiftandWidth,
     SourceGalaxySelectField,
     SourceGalaxySystematic,
     PhotoZShiftFactory,
+    PhotoZShiftandWidthFactory,
     Tracer,
 )
 from firecrown.metadata.two_point import InferredGalaxyZDist
@@ -55,6 +57,10 @@ def apply(
         """Apply method to include systematics in the tracer_arg."""
 
 
+class PhotoZShiftandWidth(SourceGalaxyPhotoZShiftandWidth[WeakLensingArgs]):
+    """Photo-z shift systematic."""
+
+
 class PhotoZShift(SourceGalaxyPhotoZShift[WeakLensingArgs]):
     """Photo-z shift systematic."""
 
@@ -433,10 +439,11 @@ def create_global(self) -> TattAlignmentSystematic:
 
 
 WeakLensingSystematicFactory = Annotated[
-    MultiplicativeShearBiasFactory
+    PhotoZShiftFactory
+    | PhotoZShiftandWidthFactory
+    | MultiplicativeShearBiasFactory
     | LinearAlignmentSystematicFactory
-    | TattAlignmentSystematicFactory
-    | PhotoZShiftFactory,
+    | TattAlignmentSystematicFactory,
     Field(discriminator="type", union_mode="left_to_right"),
 ]