From 6f9efe52a6b6b63c21b20475b09fa953b9b5b6b0 Mon Sep 17 00:00:00 2001
From: jaimerzp <jaimerz011235813@gmal.com>
Date: Fri, 19 Jul 2024 14:33:13 +0100
Subject: [PATCH] moving PhotoZShiftFactory to source

---
 firecrown/likelihood/number_counts.py | 21 +-------------------
 firecrown/likelihood/source.py        | 22 +++++++++++++++++++--
 firecrown/likelihood/weak_lensing.py  | 28 +--------------------------
 3 files changed, 22 insertions(+), 49 deletions(-)

diff --git a/firecrown/likelihood/number_counts.py b/firecrown/likelihood/number_counts.py
index 48b3a3ae..39bc96d1 100644
--- a/firecrown/likelihood/number_counts.py
+++ b/firecrown/likelihood/number_counts.py
@@ -22,6 +22,7 @@
     SourceGalaxyPhotoZShift,
     SourceGalaxySelectField,
     SourceGalaxySystematic,
+    PhotoZShiftFactory,
     Tracer,
 )
 from firecrown.metadata.two_point import InferredGalaxyZDist
@@ -476,26 +477,6 @@ def get_scale(self):
         assert self.current_tracer_args
         return self.current_tracer_args.scale
 
-
-class PhotoZShiftFactory(BaseModel):
-    """Factory class for PhotoZShift objects."""
-
-    model_config = ConfigDict(extra="forbid", frozen=True)
-
-    type: Annotated[
-        Literal["PhotoZShiftFactory"],
-        Field(description="The type of the systematic."),
-    ] = "PhotoZShiftFactory"
-
-    def create(self, bin_name: str) -> PhotoZShift:
-        """Create a PhotoZShift object with the given tracer name."""
-        return PhotoZShift(bin_name)
-
-    def create_global(self) -> PhotoZShift:
-        """Create a PhotoZShift object with the given tracer name."""
-        raise ValueError("PhotoZShift cannot be global.")
-
-
 class LinearBiasSystematicFactory(BaseModel):
     """Factory class for LinearBiasSystematic objects."""
 
diff --git a/firecrown/likelihood/source.py b/firecrown/likelihood/source.py
index 8b6c491d..22a314a6 100644
--- a/firecrown/likelihood/source.py
+++ b/firecrown/likelihood/source.py
@@ -4,8 +4,9 @@
 
 from abc import abstractmethod
 from dataclasses import dataclass, replace
-from typing import Generic, Sequence, TypeVar, final
+from typing import Generic, Sequence, TypeVar, final, Annotated, Literal
 
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
 import numpy as np
 import numpy.typing as npt
 import pyccl
@@ -32,7 +33,6 @@ class SourceSystematic(Updatable):
     def read(self, sacc_data: sacc.Sacc):
         """Call to allow this object to read from the appropriate sacc data."""
 
-
 class Source(Updatable):
     """An abstract source class (e.g., a sample of lenses).
 
@@ -248,6 +248,24 @@ def apply(self, tools: ModelingTools, tracer_arg: _SourceGalaxyArgsT):
             dndz=dndz,
         )
 
+class PhotoZShiftFactory(BaseModel):
+    """Factory class for PhotoZShift objects."""
+
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    type: Annotated[
+        Literal["PhotoZShiftFactory"],
+        Field(description="The type of the systematic."),
+    ] = "PhotoZShiftFactory"
+
+    def create(self, bin_name: str) -> PhotoZShift:
+        """Create a PhotoZShift object with the given tracer name."""
+        return PhotoZShift(bin_name)
+
+    def create_global(self) -> PhotoZShift:
+        """Create a PhotoZShift object with the given tracer name."""
+        raise ValueError("PhotoZShift cannot be global.")
+
 
 class SourceGalaxySelectField(
     SourceGalaxySystematic[_SourceGalaxyArgsT], Generic[_SourceGalaxyArgsT]
diff --git a/firecrown/likelihood/weak_lensing.py b/firecrown/likelihood/weak_lensing.py
index b1102df8..a92fbc95 100644
--- a/firecrown/likelihood/weak_lensing.py
+++ b/firecrown/likelihood/weak_lensing.py
@@ -23,6 +23,7 @@
     SourceGalaxyPhotoZShift,
     SourceGalaxySelectField,
     SourceGalaxySystematic,
+    PhotoZShiftFactory,
     Tracer,
 )
 from firecrown.metadata.two_point import InferredGalaxyZDist
@@ -430,33 +431,6 @@ def create_global(self) -> TattAlignmentSystematic:
         """
         return TattAlignmentSystematic(None)
 
-
-class PhotoZShiftFactory(BaseModel):
-    """Factory class for PhotoZShift objects."""
-
-    model_config = ConfigDict(extra="forbid", frozen=True)
-
-    type: Annotated[
-        Literal["PhotoZShiftFactory"], Field(description="The type of the systematic.")
-    ] = "PhotoZShiftFactory"
-
-    def create(self, bin_name: str) -> PhotoZShift:
-        """Create a PhotoZShift object.
-
-        :param inferred_zdist: The inferred galaxy redshift distribution for
-            the created PhotoZShift object.
-        :return: The created PhotoZShift object.
-        """
-        return PhotoZShift(bin_name)
-
-    def create_global(self) -> PhotoZShift:
-        """Create a PhotoZShift object.
-
-        :return: The created PhotoZShift object.
-        """
-        raise ValueError("PhotoZShift cannot be global")
-
-
 WeakLensingSystematicFactory = Annotated[
     MultiplicativeShearBiasFactory
     | LinearAlignmentSystematicFactory