Merge pull request #37 from Ciela-Institute/simulatorclass

Added Simulator class object

src/caustic/lenses/base.py

@@ -1,5 +1,6 @@
 from abc import abstractmethod
 from typing import Any, Optional
+from functools import partial
 
 import torch
 from torch import Tensor
@@ -11,7 +12,6 @@
 
 __all__ = ("ThinLens", "ThickLens")
 
-
 class ThickLens(Parametrized):
     """
     Base class for modeling gravitational lenses that cannot be treated using the thin lens approximation.
@@ -117,8 +117,7 @@ def magnification(self, thx: Tensor, thy: Tensor, z_s: Tensor, x) -> Tensor:
         Returns:
             Tensor: The gravitational lensing magnification at the given coordinates.
         """
-        return get_magnification(self.raytrace, thx, thy, z_s, x)
-
+        return get_magnification(partial(self.raytrace, x = x), thx, thy, z_s)
 
 class ThinLens(Parametrized):
     """Base class for thin gravitational lenses.
@@ -328,4 +327,4 @@ def magnification(self, thx: Tensor, thy: Tensor, z_s: Tensor, x) -> Tensor:
         Returns:
             Tensor: Gravitational magnification at the given coordinates.
         """
-        return get_magnification(self.raytrace, thx, thy, z_s, x)
+        return get_magnification(partial(self.raytrace, x = x), thx, thy, z_s)

src/caustic/lenses/utils.py

@@ -9,7 +9,7 @@
 
 
 def get_pix_jacobian(
-    raytrace, thx, thy, z_s, x
+    raytrace, thx, thy, z_s
 ) -> Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor]]:
     """Computes the Jacobian matrix of the partial derivatives of the
     image position with respect to the source position
@@ -27,11 +27,11 @@ def get_pix_jacobian(
         The Jacobian matrix of the image position with respect to the source position at the given point.
 
     """
-    jac = torch.func.jacfwd(raytrace, (0, 1))(thx, thy, z_s, x)  # type: ignore
+    jac = torch.func.jacfwd(raytrace, (0, 1))(thx, thy, z_s)  # type: ignore
     return jac
 
 
-def get_pix_magnification(raytrace, thx, thy, z_s, x) -> Tensor:
+def get_pix_magnification(raytrace, thx, thy, z_s) -> Tensor:
     """
     Computes the magnification at a single point on the lensing plane. The magnification is derived from the determinant
     of the Jacobian matrix of the image position with respect to the source position.
@@ -46,11 +46,11 @@ def get_pix_magnification(raytrace, thx, thy, z_s, x) -> Tensor:
     Returns:
         The magnification at the given point on the lensing plane.
     """
-    jac = get_pix_jacobian(raytrace, thx, thy, z_s, x)
+    jac = get_pix_jacobian(raytrace, thx, thy, z_s)
     return 1 / (jac[0][0] * jac[1][1] - jac[0][1] * jac[1][0]).abs()
 
 
-def get_magnification(raytrace, thx, thy, z_s, x) -> Tensor:
+def get_magnification(raytrace, thx, thy, z_s) -> Tensor:
     """
     Computes the magnification over a grid on the lensing plane. This is done by calling `get_pix_magnification` 
     for each point on the grid.
@@ -65,6 +65,6 @@ def get_magnification(raytrace, thx, thy, z_s, x) -> Tensor:
     Returns:
         A tensor representing the magnification at each point on the grid.
     """
-    return vmap_n(get_pix_magnification, 2, (None, 0, 0, None, None))(
-        raytrace, thx, thy, z_s, x
+    return vmap_n(get_pix_magnification, 2, (None, 0, 0, None))(
+        raytrace, thx, thy, z_s
     )

src/caustic/packed.py

@@ -1,6 +1,5 @@
 from collections import OrderedDict
 
-
 class Packed(OrderedDict):
     """
     Dummy wrapper for `x` so other functions can check its type.

src/caustic/parametrized.py

@@ -12,7 +12,6 @@
 
 __all__ = ("Parametrized",)
 
-
 class Parametrized:
     """
     Represents a class with Param and Parametrized attributes, typically used to construct parts of a simulator
@@ -282,7 +281,7 @@ def pack(
             ValueError: If the number of dynamic arguments does not match the expected number.
             ValueError: If the input is a tensor and the shape does not match the expected shape.
         """
-        if isinstance(x, dict):
+        if isinstance(x, (dict, Packed)):
             missing_names = [
                 name for name in chain([self.name], self._descendants) if name not in x
             ]
@@ -302,7 +301,7 @@ def pack(
                 # TODO: give component and arg names
                 raise ValueError(
                     f"{n_passed} dynamic args were passed, but {n_expected} are "
-                    "required"
+                    "required."
                 )
 
             cur_offset = self.n_dynamic
@@ -559,8 +558,7 @@ def add_params(p: Parametrized, dot):
                 add_params(desc, dot)
 
         return dot
-
-
+
 # class ParametrizedList(Parametrized):
 #     """
 #     TODO

src/caustic/simulator.py

@@ -0,0 +1,18 @@
+from .parameterized import Parametrized
+
+__all__ = ("Simulator", )
+
+class Simulator(Parametrized):
+    """A caustic simulator using Parametrized framework.
+
+    Defines a simulator class which is a callable function that
+    operates on the Parametrized framework. Users define the `forward`
+    method which takes as its first argument an object which can be
+    packed, all other args and kwargs are simply passed to the forward
+    method.
+
+    See `Parametrized` for details on how to add/access parameters.
+
+    """
+    def __call__(self, *args, **kwargs):
+        return self.forward(self.pack(args[0]), *args[1:], **kwargs)