Merge pull request #38 from Ciela-Institute/documentation

Adding documentation to all classes and methods in google style

src/caustic/cosmology.py

@@ -37,43 +37,131 @@
 
 class Cosmology(Parametrized):
     """
+    Abstract base class for cosmological models.
+
+    This class provides an interface for cosmological computations used in lensing 
+    such as comoving distance and critical surface density. 
+
     Units:
         - Distance: Mpc
-        - Mass: solMass
+        - Mass: solar mass
+
+    Attributes:
+        name (str): Name of the cosmological model.
     """
 
     def __init__(self, name: str):
+        """
+        Initialize the Cosmology.
+
+        Args:
+            name (str): Name of the cosmological model.
+        """
         super().__init__(name)
 
     @abstractmethod
     def rho_cr(self, z: Tensor, x: Optional[dict[str, Any]] = None) -> Tensor:
+        """
+        Compute the critical density at redshift z.
+
+        Args:
+            z (Tensor): The redshifts.
+            x (Optional[dict[str, Any]]): Additional parameters for the computation.
+
+        Returns:
+            Tensor: The critical density at each redshift.
+        """
         ...
 
     @abstractmethod
     def comoving_dist(self, z: Tensor, x: Optional[dict[str, Any]] = None) -> Tensor:
+        """
+        Compute the comoving distance to redshift z.
+
+        Args:
+            z (Tensor): The redshifts.
+            x (Optional[dict[str, Any]]): Additional parameters for the computation.
+
+        Returns:
+            Tensor: The comoving distance to each redshift.
+        """
         ...
 
     def comoving_dist_z1z2(
         self, z1: Tensor, z2: Tensor, x: Optional[dict[str, Any]] = None) -> Tensor:
+        """
+        Compute the comoving distance between two redshifts.
+
+        Args:
+            z1 (Tensor): The starting redshifts.
+            z2 (Tensor): The ending redshifts.
+            x (Optional[dict[str, Any]]): Additional parameters for the computation.
+
+        Returns:
+            Tensor: The comoving distance between each pair of redshifts.
+        """
         return self.comoving_dist(z2, x) - self.comoving_dist(z1, x)
 
     def angular_diameter_dist(
         self, z: Tensor, x: Optional[dict[str, Any]] = None) -> Tensor:
+        """
+        Compute the angular diameter distance to redshift z.
+
+        Args:
+            z (Tensor): The redshifts.
+            x (Optional[dict[str, Any]]): Additional parameters for the computation.
+
+        Returns:
+            Tensor: The angular diameter distance to each redshift.
+        """
         return self.comoving_dist(z, x) / (1 + z)
 
     def angular_diameter_dist_z1z2(
         self, z1: Tensor, z2: Tensor, x: Optional[dict[str, Any]] = None) -> Tensor:
+        """
+        Compute the angular diameter distance between two redshifts.
+
+        Args:
+            z1 (Tensor): The starting redshifts.
+            z2 (Tensor): The ending redshifts.
+            x (Optional[dict[str, Any]]): Additional parameters for the computation.
+
+        Returns:
+            Tensor: The angular diameter distance between each pair of redshifts.
+        """
         return self.comoving_dist_z1z2(z1, z2, x) / (1 + z2)
 
     def time_delay_dist(
         self, z_l: Tensor, z_s: Tensor, x: Optional[dict[str, Any]] = None) -> Tensor:
+        """
+        Compute the time delay distance between lens and source planes.
+
+        Args:
+            z_l (Tensor): The lens redshifts.
+            z_s (Tensor): The source redshifts.
+            x (Optional[dict[str, Any]]): Additional parameters for the computation.
+
+        Returns:
+            Tensor: The time delay distance for each pair of lens and source redshifts.
+        """
         d_l = self.angular_diameter_dist(z_l, x)
         d_s = self.angular_diameter_dist(z_s, x)
         d_ls = self.angular_diameter_dist_z1z2(z_l, z_s, x)
         return (1 + z_l) * d_l * d_s / d_ls
 
     def Sigma_cr(
         self, z_l: Tensor, z_s: Tensor, x: Optional[dict[str, Any]] = None) -> Tensor:
+        """
+        Compute the critical surface density between lens and source planes.
+
+        Args:
+            z_l (Tensor): The lens redshifts.
+            z_s (Tensor): The source redshifts.
+            x (Optional[dict[str, Any]]): Additional parameters for the computation.
+
+        Returns:
+            Tensor: The critical surface density for each pair of lens and source redshifts.
+        """
         d_l = self.angular_diameter_dist(z_l, x)
         d_s = self.angular_diameter_dist(z_s, x)
         d_ls = self.angular_diameter_dist_z1z2(z_l, z_s, x)
@@ -82,17 +170,25 @@ def Sigma_cr(
 
 class FlatLambdaCDM(Cosmology):
     """
-
-    Flat LCDM cosmology with no radiation.
+    Subclass of Cosmology representing a Flat Lambda Cold Dark Matter (LCDM) cosmology with no radiation.
     """
-
+    
     def __init__(
         self,
         name: str,
         h0: Optional[Tensor] = torch.tensor(h0_default),
         rho_cr_0: Optional[Tensor] = torch.tensor(rho_cr_0_default),
         Om0: Optional[Tensor] = torch.tensor(Om0_default),
     ):
+        """
+        Initialize a new instance of the FlatLambdaCDM class.
+
+        Args:
+            name (str): Name of the cosmology.
+            h0 (Optional[Tensor]): Hubble constant. Default is h0_default.
+            rho_cr_0 (Optional[Tensor]): Critical density at z=0. Default is rho_cr_0_default.
+            Om0 (Optional[Tensor]): Matter density parameter at z=0. Default is Om0_default.
+        """
         super().__init__(name)
 
         self.add_param("h0", h0)
@@ -107,21 +203,59 @@ def __init__(
         )
 
     def dist_hubble(self, h0):
+        """
+        Calculate the Hubble distance.
+
+        Args:
+            h0 (Tensor): Hubble constant.
+
+        Returns:
+            Tensor: Hubble distance.
+        """
         return c_Mpc_s / (100 * km_to_Mpc) / h0
 
     def rho_cr(self, z: Tensor, x: Optional[dict[str, Any]] = None) -> torch.Tensor:
+        """
+        Calculate the critical density at redshift z.
+
+        Args:
+            z (Tensor): Redshift.
+            x (Optional[dict[str, Any]]): Additional parameters for the computation.
+
+        Returns:
+            torch.Tensor: Critical density at redshift z.
+        """
         _, rho_cr_0, Om0 = self.unpack(x)
         Ode0 = 1 - Om0
         return rho_cr_0 * (Om0 * (1 + z) ** 3 + Ode0)
 
     def _comoving_dist_helper(self, x: Tensor) -> Tensor:
+        """
+        Helper method for computing comoving distances.
+
+        Args:
+            x (Tensor): Input tensor.
+
+        Returns:
+            Tensor: Computed comoving distances.
+        """
         return interp1d(
             self._comoving_dist_helper_x_grid,
             self._comoving_dist_helper_y_grid,
             torch.atleast_1d(x),
         ).reshape(x.shape)
 
     def comoving_dist(self, z: Tensor, x: Optional[dict[str, Any]] = None) -> Tensor:
+        """
+        Calculate the comoving distance to redshift z.
+
+        Args:
+            z (Tensor): Redshift.
+            x (Optional[dict[str, Any]]): Additional parameters for the computation.
+
+        Returns:
+            Tensor: Comoving distance to redshift z.
+        """
         h0, _, Om0 = self.unpack(x)
 
         Ode0 = 1 - Om0

src/caustic/fwd_raytrace.py

@@ -19,6 +19,38 @@ def fwd_raytrace(
     max_iters_guesses: int = 50,
     max_iters_final: int = 50,
 ) -> Tensor:
+    """
+    Implements a forward ray tracing algorithm for a strong gravitational lensing system. 
+
+    Args:
+        beta_x (Tensor): The x coordinates of the source positions in the source plane.
+        beta_y (Tensor): The y coordinates of the source positions in the source plane.
+        get_beta_hat (Callable[[Tensor, Tensor], Tuple[Tensor, Tensor]]): A function that returns 
+            the predicted source positions given the lensed image positions.
+        n_images (int): The number of images to produce.
+        thx_range (Tuple[float, float]): The range of x coordinates in the lens plane to consider for initial guesses.
+        thy_range (Tuple[float, float]): The range of y coordinates in the lens plane to consider for initial guesses.
+        n_guesses (int, optional): The number of initial guesses for the lensed image positions. Default is 15.
+        lam (float, optional): The damping parameter for the Levenberg-Marquardt optimization. Default is 1e-2.
+        max_iters_guesses (int, optional): The maximum number of iterations for the optimization of initial guesses. Default is 50.
+        max_iters_final (int, optional): The maximum number of iterations for the final optimization. Default is 50.
+
+    Returns:
+        Tensor: The optimized lensed image positions in the lens plane.
+    
+    This function first generates a set of initial guesses for the lensed image positions. 
+    These guesses are then optimized using the Levenberg-Marquardt algorithm to match the 
+    observed source positions. If the optimization fails for any of the initial guesses, 
+    new guesses are generated and the process is repeated until a successful optimization is achieved.
+
+    Once the initial optimization is complete, the results are pared down to the desired number 
+    of images using a clustering algorithm, and a final round of optimization is performed. 
+    The function returns the final optimized lensed image positions.
+
+    Note: If the number of images is greater than the number of observed source positions, 
+    the function may not be able to find a solution.
+    """
+
     bxy = torch.stack((beta_x, beta_y))
     thxy_min = torch.tensor((thx_range[0], thy_range[0]))
     thxy_max = torch.tensor((thx_range[1], thy_range[1]))