From d7363ccd60411733e39e6351ca42c4da5f5510f1 Mon Sep 17 00:00:00 2001
From: jaimerzp <jaimerz011235813@gmal.com>
Date: Fri, 19 Jul 2024 17:32:55 +0100
Subject: [PATCH] widths in shidts in WL and NC

---
 firecrown/likelihood/number_counts.py |  7 +++++++
 firecrown/likelihood/weak_lensing.py  | 13 ++++++++++---
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/firecrown/likelihood/number_counts.py b/firecrown/likelihood/number_counts.py
index 34adb6ea..199e451e 100644
--- 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
index 82a12a71..a36c218f 100644
--- 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"),
 ]