From 5248a82fa3ecdfff3b73046e4f32033e4e3b7e3f Mon Sep 17 00:00:00 2001
From: Tilman Troester <tilman.troester@gmail.com>
Date: Thu, 3 Jul 2025 18:48:45 +0200
Subject: [PATCH] added SourceGalaxySelectFieldFactory

---
 firecrown/likelihood/number_counts.py |  4 +++-
 firecrown/likelihood/source.py        | 21 +++++++++++++++++++++
 firecrown/likelihood/weak_lensing.py  |  4 +++-
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/firecrown/likelihood/number_counts.py b/firecrown/likelihood/number_counts.py
index b2c2fd14c..d15afa59b 100644
--- a/firecrown/likelihood/number_counts.py
+++ b/firecrown/likelihood/number_counts.py
@@ -22,6 +22,7 @@
     SourceGalaxySystematic,
     PhotoZShiftFactory,
     PhotoZShiftandStretchFactory,
+    SourceGalaxySelectFieldFactory,
     Tracer,
 )
 from firecrown.metadata_types import InferredGalaxyZDist, TypeSource
@@ -613,7 +614,8 @@ def create_global(self) -> ConstantMagnificationBiasSystematic:
     | LinearBiasSystematicFactory
     | PTNonLinearBiasSystematicFactory
     | MagnificationBiasSystematicFactory
-    | ConstantMagnificationBiasSystematicFactory,
+    | ConstantMagnificationBiasSystematicFactory
+    | SourceGalaxySelectFieldFactory,
     Field(discriminator="type", union_mode="left_to_right"),
 ]
 
diff --git a/firecrown/likelihood/source.py b/firecrown/likelihood/source.py
index a0c87ab1f..a67b8c75a 100644
--- a/firecrown/likelihood/source.py
+++ b/firecrown/likelihood/source.py
@@ -476,6 +476,27 @@ def apply(
         return replace(tracer_arg, field=self.field)
 
 
+class SourceGalaxySelectFieldFactory(BaseModel):
+    """Factory class for SourceGalaxySelectField objects."""
+
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    field_name: str
+
+    type: Annotated[
+        Literal["SourceGalaxySelectFieldFactory"],
+        Field(description="The type of the systematic."),
+    ] = "SourceGalaxySelectFieldFactory"
+
+    def create(self, bin_name: str) -> SourceGalaxySelectField:
+        """Create a SourceGalaxySelectField object with the given field name."""
+        raise ValueError("SourceGalaxySelectField cannot be per-bin.")
+
+    def create_global(self) -> SourceGalaxySelectField:
+        """Create a PhotoZShift object with the given field name."""
+        return SourceGalaxySelectField(field=self.field_name)
+
+
 class SourceGalaxy(Source, Generic[_SourceGalaxyArgsT]):
     """Source class for galaxy based sources."""
 
diff --git a/firecrown/likelihood/weak_lensing.py b/firecrown/likelihood/weak_lensing.py
index 3141ded6f..2b58a53ae 100644
--- a/firecrown/likelihood/weak_lensing.py
+++ b/firecrown/likelihood/weak_lensing.py
@@ -24,6 +24,7 @@
     SourceGalaxySystematic,
     PhotoZShiftFactory,
     PhotoZShiftandStretchFactory,
+    SourceGalaxySelectFieldFactory,
     Tracer,
 )
 from firecrown.metadata_types import InferredGalaxyZDist, TypeSource
@@ -516,7 +517,8 @@ def create_global(self) -> TattAlignmentSystematic:
     | PhotoZShiftandStretchFactory
     | MultiplicativeShearBiasFactory
     | LinearAlignmentSystematicFactory
-    | TattAlignmentSystematicFactory,
+    | TattAlignmentSystematicFactory
+    | SourceGalaxySelectFieldFactory,
     Field(discriminator="type", union_mode="left_to_right"),
 ]