feat: Add functionality to build simulator from a YAML configuration …

…file (Ciela-Institute#167)

* feat: Add registry for various parametrized "kind" (#84)

* feat: Add registry for various parametrized "kind"

Added a registry for cosmology, lenses, light,
and sims classes to be used as "kind"

* style: pre-commit fixes

* fix: Fix misspelling of pixelatedconvergence

Co-authored-by: Cordero Core <127983572+uwcdc@users.noreply.github.com>

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Cordero Core <127983572+uwcdc@users.noreply.github.com>

* feat: Add _meta_params class attribute

* fix: Fix default _meta_params type

* feat: Add _meta_params to cosmo, lenses, and source (#87)

* feat: Add _meta_params to cosmology, lenses, and sources for use in yaml validation

* feat: Add _meta_params to lens_source

* fix: Remove _meta_params from multiplane.py

Co-authored-by: Don Setiawan <landungs@uw.edu>

* fix: Remove _meta_params from singleplane.py

* fix: Update _meta_params to reflect z_s

---------

Co-authored-by: Don Setiawan <landungs@uw.edu>

* feat: Add pydantic dynamic models from classes (#88)

* feat: Add pydantic dynamic models from classes

* revert: Remove the use of _meta_params, and use annotated instead

* feat: Updated dynamic creation to use annotated from class

* fix: Fix bugs in creating field defs

* test: Add tests for models/utils

* fix: Fix typehints for Parametrized |

* chore(deps): Add pydantic 2 as dependency

* refactor: Add build_simulator function to caustics init

* feat: Add way to evaluate string and dict in pre field inputs

* fix: Fix return type to be Any

* fix: Apply suggestions from code review

Co-authored-by: Cordero Core <127983572+uwcdc@users.noreply.github.com>

* fix: Apply suggestions from code review

Co-authored-by: Cordero Core <127983572+uwcdc@users.noreply.github.com>

---------

Co-authored-by: Cordero Core <127983572+uwcdc@users.noreply.github.com>

* feat: Add ValueError when dict doesn't include both 'func' and 'keys'

* test: Add integration test for yaml config (#89)

* test: Add integration test for yaml config

* refactor: Remove unecessary attr and specify union for build

* test: Add test for models registry

* chore(deps): Add pytest-mock for mocking

* test: Separate models test utils and create a complex yaml

* fix: Move SinglePlane to single lenses and handle case

* test: Add test for models api

* fix: Add single plane model to list of lenses

* refactor: Renamed config_json to config_dict for clarity

* test: Fix tempfile test for windows

* test: Fix temp file creation and reading

* test: Extract temp yaml creation to a func

* test: Fix to use path of tempfile

* test: Fix where name is extracted

* test: Add build_simulator test with state

* test: Ignore cleanup for temp state dict... let OS clean it up

* fix: Fix bug with arbitrary dict

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Cordero Core <127983572+uwcdc@users.noreply.github.com>

pyproject.toml

@@ -47,7 +47,8 @@ dev = [
     "lenstronomy==1.11.1",
     "pytest>=8.0,<9",
     "pytest-cov>=4.1,<5",
-    "pre-commit>=3.6,<4"
+    "pytest-mock>=3.12,<4",
+    "pre-commit>=3.6,<4",
 ]
 
 [tool.hatch.metadata.hooks.requirements_txt]

requirements.txt

@@ -2,6 +2,7 @@ astropy>=5.2.1,<6.0.0
 graphviz==0.20.1
 h5py>=3.8.0
 numpy>=1.23.5
+pydantic>=2.6.1,<3
 safetensors>=0.4.1
 scipy>=1.8.0
 torch>=2.0.0

src/caustics/__init__.py

@@ -28,6 +28,7 @@
 from . import utils
 from .sims import Lens_Source, Simulator
 from .tests import test
+from .models.api import build_simulator
 
 __version__ = VERSION
 __author__ = "Ciela"
@@ -62,4 +63,5 @@
     "Lens_Source",
     "Simulator",
     "test",
+    "build_simulator",
 ]

src/caustics/cosmology/FlatLambdaCDM.py

@@ -1,5 +1,5 @@
 # mypy: disable-error-code="operator"
-from typing import Optional
+from typing import Optional, Annotated
 
 import torch
 from torch import Tensor
@@ -11,9 +11,7 @@
 from ..parametrized import unpack
 from ..packed import Packed
 from ..constants import c_Mpc_s, km_to_Mpc
-from .base import (
-    Cosmology,
-)
+from .base import Cosmology, NameType
 
 _h0_default = float(default_cosmology.get().h)
 _critical_density_0_default = float(
@@ -43,10 +41,14 @@ class FlatLambdaCDM(Cosmology):
 
     def __init__(
         self,
-        h0: Optional[Tensor] = h0_default,
-        critical_density_0: Optional[Tensor] = critical_density_0_default,
-        Om0: Optional[Tensor] = Om0_default,
-        name: Optional[str] = None,
+        h0: Annotated[Optional[Tensor], "Hubble constant over 100", True] = h0_default,
+        critical_density_0: Annotated[
+            Optional[Tensor], "Critical density at z=0", True
+        ] = critical_density_0_default,
+        Om0: Annotated[
+            Optional[Tensor], "Matter density parameter at z=0", True
+        ] = Om0_default,
+        name: NameType = None,
     ):
         """
         Initialize a new instance of the FlatLambdaCDM class.

src/caustics/cosmology/base.py

@@ -1,14 +1,16 @@
 # mypy: disable-error-code="operator"
 from abc import abstractmethod
 from math import pi
-from typing import Optional
+from typing import Optional, Annotated
 
 from torch import Tensor
 
 from ..constants import G_over_c2
 from ..parametrized import Parametrized, unpack
 from ..packed import Packed
 
+NameType = Annotated[Optional[str], "Name of the cosmology"]
+
 
 class Cosmology(Parametrized):
     """
@@ -31,7 +33,7 @@ class Cosmology(Parametrized):
         Name of the cosmological model.
     """
 
-    def __init__(self, name: Optional[str] = None):
+    def __init__(self, name: NameType = None):
         """
         Initialize the Cosmology.
 

src/caustics/lenses/base.py

@@ -1,6 +1,6 @@
 # mypy: disable-error-code="call-overload"
 from abc import abstractmethod
-from typing import Optional, Union
+from typing import Optional, Union, Annotated, List
 from functools import partial
 import warnings
 
@@ -16,13 +16,21 @@
 
 __all__ = ("ThinLens", "ThickLens")
 
+CosmologyType = Annotated[
+    Cosmology,
+    "Cosmology object that encapsulates cosmological parameters and distances",
+]
+NameType = Annotated[Optional[str], "Name of the lens model"]
+ZLType = Annotated[Optional[Union[Tensor, float]], "The redshift of the lens", True]
+LensesType = Annotated[List["ThinLens"], "A list of ThinLens objects"]
+
 
 class Lens(Parametrized):
     """
     Base class for all lenses
     """
 
-    def __init__(self, cosmology: Cosmology, name: Optional[str] = None):
+    def __init__(self, cosmology: CosmologyType, name: NameType = None):
         """
         Initializes a new instance of the Lens class.
 
@@ -715,9 +723,9 @@ class ThinLens(Lens):
 
     def __init__(
         self,
-        cosmology: Cosmology,
-        z_l: Optional[Union[Tensor, float]] = None,
-        name: Optional[str] = None,
+        cosmology: CosmologyType,
+        z_l: ZLType = None,
+        name: NameType = None,
     ):
         super().__init__(cosmology=cosmology, name=name)
         self.add_param("z_l", z_l)

src/caustics/lenses/epl.py

@@ -1,12 +1,11 @@
-# mypy: disable-error-code="operator"
-from typing import Optional, Union
+# mypy: disable-error-code="operator,dict-item"
+from typing import Optional, Union, Annotated
 
 import torch
 from torch import Tensor
 
-from ..cosmology import Cosmology
 from ..utils import derotate, translate_rotate
-from .base import ThinLens
+from .base import ThinLens, CosmologyType, NameType, ZLType
 from ..parametrized import unpack
 from ..packed import Packed
 
@@ -92,17 +91,33 @@ class EPL(ThinLens):
 
     def __init__(
         self,
-        cosmology: Cosmology,
-        z_l: Optional[Union[Tensor, float]] = None,
-        x0: Optional[Union[Tensor, float]] = None,
-        y0: Optional[Union[Tensor, float]] = None,
-        q: Optional[Union[Tensor, float]] = None,
-        phi: Optional[Union[Tensor, float]] = None,
-        b: Optional[Union[Tensor, float]] = None,
-        t: Optional[Union[Tensor, float]] = None,
-        s: float = 0.0,
-        n_iter: int = 18,
-        name: Optional[str] = None,
+        cosmology: CosmologyType,
+        z_l: ZLType = None,
+        x0: Annotated[
+            Optional[Union[Tensor, float]], "X coordinate of the lens center", True
+        ] = None,
+        y0: Annotated[
+            Optional[Union[Tensor, float]], "Y coordinate of the lens center", True
+        ] = None,
+        q: Annotated[
+            Optional[Union[Tensor, float]], "Axis ratio of the lens", True
+        ] = None,
+        phi: Annotated[
+            Optional[Union[Tensor, float]], "Position angle of the lens", True
+        ] = None,
+        b: Annotated[
+            Optional[Union[Tensor, float]], "Scale length of the lens", True
+        ] = None,
+        t: Annotated[
+            Optional[Union[Tensor, float]],
+            "Power law slope (`gamma-1`) of the lens",
+            True,
+        ] = None,
+        s: Annotated[
+            float, "Softening length for the elliptical power-law profile"
+        ] = 0.0,
+        n_iter: Annotated[int, "Number of iterations for the iterative solver"] = 18,
+        name: NameType = None,
     ):
         """
         Initialize an EPL lens model.

src/caustics/lenses/external_shear.py

@@ -1,10 +1,10 @@
-from typing import Optional, Union
+# mypy: disable-error-code="dict-item"
+from typing import Optional, Union, Annotated
 
 from torch import Tensor
 
-from ..cosmology import Cosmology
 from ..utils import translate_rotate
-from .base import ThinLens
+from .base import ThinLens, CosmologyType, NameType, ZLType
 from ..parametrized import unpack
 from ..packed import Packed
 
@@ -53,14 +53,28 @@ class ExternalShear(ThinLens):
 
     def __init__(
         self,
-        cosmology: Cosmology,
-        z_l: Optional[Union[Tensor, float]] = None,
-        x0: Optional[Union[Tensor, float]] = None,
-        y0: Optional[Union[Tensor, float]] = None,
-        gamma_1: Optional[Union[Tensor, float]] = None,
-        gamma_2: Optional[Union[Tensor, float]] = None,
-        s: float = 0.0,
-        name: Optional[str] = None,
+        cosmology: CosmologyType,
+        z_l: ZLType = None,
+        x0: Annotated[
+            Optional[Union[Tensor, float]],
+            "x-coordinate of the shear center in the lens plane",
+            True,
+        ] = None,
+        y0: Annotated[
+            Optional[Union[Tensor, float]],
+            "y-coordinate of the shear center in the lens plane",
+            True,
+        ] = None,
+        gamma_1: Annotated[
+            Optional[Union[Tensor, float]], "Shear component in the x-direction", True
+        ] = None,
+        gamma_2: Annotated[
+            Optional[Union[Tensor, float]], "Shear component in the y-direction", True
+        ] = None,
+        s: Annotated[
+            float, "Softening length for the elliptical power-law profile"
+        ] = 0.0,
+        name: NameType = None,
     ):
         super().__init__(cosmology, z_l, name=name)
 

src/caustics/lenses/mass_sheet.py

@@ -1,12 +1,11 @@
-# mypy: disable-error-code="operator"
-from typing import Optional, Union
+# mypy: disable-error-code="operator,dict-item"
+from typing import Optional, Union, Annotated
 
 import torch
 from torch import Tensor
 
-from ..cosmology import Cosmology
 from ..utils import translate_rotate
-from .base import ThinLens
+from .base import ThinLens, CosmologyType, NameType, ZLType
 from ..parametrized import unpack
 from ..packed import Packed
 
@@ -64,12 +63,22 @@ class MassSheet(ThinLens):
 
     def __init__(
         self,
-        cosmology: Cosmology,
-        z_l: Optional[Union[Tensor, float]] = None,
-        x0: Optional[Union[Tensor, float]] = None,
-        y0: Optional[Union[Tensor, float]] = None,
-        surface_density: Optional[Union[Tensor, float]] = None,
-        name: Optional[str] = None,
+        cosmology: CosmologyType,
+        z_l: ZLType = None,
+        x0: Annotated[
+            Optional[Union[Tensor, float]],
+            "x-coordinate of the shear center in the lens plane",
+            True,
+        ] = None,
+        y0: Annotated[
+            Optional[Union[Tensor, float]],
+            "y-coordinate of the shear center in the lens plane",
+            True,
+        ] = None,
+        surface_density: Annotated[
+            Optional[Union[Tensor, float]], "Surface density", True
+        ] = None,
+        name: NameType = None,
     ):
         super().__init__(cosmology, z_l, name=name)
 

src/caustics/lenses/multiplane.py

@@ -5,8 +5,7 @@
 from torch import Tensor
 
 from ..constants import arcsec_to_rad, rad_to_arcsec, c_Mpc_s
-from ..cosmology import Cosmology
-from .base import ThickLens, ThinLens
+from .base import ThickLens, NameType, CosmologyType, LensesType
 from ..parametrized import unpack
 from ..packed import Packed
 
@@ -19,7 +18,7 @@ class Multiplane(ThickLens):
 
     Attributes
     ----------
-    lenses (list[ThinLens])
+    lenses list of ThinLens
         List of thin lenses.
 
     Parameters
@@ -33,7 +32,7 @@ class Multiplane(ThickLens):
     """
 
     def __init__(
-        self, cosmology: Cosmology, lenses: list[ThinLens], name: Optional[str] = None
+        self, cosmology: CosmologyType, lenses: LensesType, name: NameType = None
     ):
         super().__init__(cosmology, name=name)
         self.lenses = lenses

src/caustics/lenses/nfw.py

@@ -1,14 +1,13 @@
-# mypy: disable-error-code="operator,union-attr"
+# mypy: disable-error-code="operator,union-attr,dict-item"
 from math import pi
-from typing import Optional, Union
+from typing import Optional, Union, Annotated, Literal
 
 import torch
 from torch import Tensor
 
 from ..constants import G_over_c2, arcsec_to_rad, rad_to_arcsec
-from ..cosmology import Cosmology
 from ..utils import translate_rotate
-from .base import ThinLens
+from .base import ThinLens, NameType, CosmologyType, ZLType
 from ..parametrized import unpack
 from ..packed import Packed
 
@@ -100,15 +99,26 @@ class NFW(ThinLens):
 
     def __init__(
         self,
-        cosmology: Cosmology,
-        z_l: Optional[Union[Tensor, float]] = None,
-        x0: Optional[Union[Tensor, float]] = None,
-        y0: Optional[Union[Tensor, float]] = None,
-        m: Optional[Union[Tensor, float]] = None,
-        c: Optional[Union[Tensor, float]] = None,
-        s: float = 0.0,
-        use_case="batchable",
-        name: Optional[str] = None,
+        cosmology: CosmologyType,
+        z_l: ZLType = None,
+        x0: Annotated[
+            Optional[Union[Tensor, float]], "X coordinate of the lens center", True
+        ] = None,
+        y0: Annotated[
+            Optional[Union[Tensor, float]], "Y coordinate of the lens center", True
+        ] = None,
+        m: Annotated[Optional[Union[Tensor, float]], "Mass of the lens", True] = None,
+        c: Annotated[
+            Optional[Union[Tensor, float]], "Concentration parameter of the lens", True
+        ] = None,
+        s: Annotated[
+            float,
+            "Softening parameter to avoid singularities at the center of the lens",
+        ] = 0.0,
+        use_case: Annotated[
+            Literal["batchable", "differentiable"], "the NFW/TNFW profile"
+        ] = "batchable",
+        name: NameType = None,
     ):
         """
         Initialize an instance of the NFW lens class.