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