Added units to the docstrings

src/caustics/cosmology.py

@@ -43,12 +43,13 @@ class Cosmology(Parametrized):
     This class provides an interface for cosmological computations used in lensing
     such as comoving distance and critical surface density.
 
-    Units
-    -----
     Distance
-        Mpc
+
+        *Unit: megaparsec*
+
     Mass
-        solar mass
+
+        *Unit: solMass*
 
     Attributes
     ----------
@@ -76,13 +77,19 @@ def critical_density(self, z: Tensor, params: Optional["Packed"] = None) -> Tens
         ----------
         z: Tensor
             The redshifts.
+
+            *Unit: unitless*
+
         params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The critical density at each redshift.
+
+            *Unit: solMass/megaparsec^3*
+
         """
         ...
 
@@ -98,13 +105,19 @@ def comoving_distance(
         ----------
         z: Tensor
             The redshifts.
-        params: (Packed, optional0
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The comoving distance to each redshift.
+
+            *Unit: megaparsec*
+
         """
         ...
 
@@ -120,13 +133,19 @@ def transverse_comoving_distance(
         ----------
         z: Tensor
             The redshifts.
+
+            *Unit: unitless*
+
         params: (Packed, optional)
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The transverse comoving distance to each redshift in Mpc.
+
+            *Unit: megaparsec*
+
         """
         ...
 
@@ -141,15 +160,24 @@ def comoving_distance_z1z2(
         ----------
         z1: Tensor
             The starting redshifts.
+
+            *Unit: unitless*
+
         z2: Tensor
             The ending redshifts.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The comoving distance between each pair of redshifts.
+
+            *Unit: megaparsec*
+
         """
         return self.comoving_distance(z2, params) - self.comoving_distance(z1, params)
 
@@ -164,15 +192,24 @@ def transverse_comoving_distance_z1z2(
         ----------
         z1: Tensor
             The starting redshifts.
+
+            *Unit: unitless*
+
         z2: Tensor
             The ending redshifts.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The transverse comoving distance between each pair of redshifts in Mpc.
+
+            *Unit: megaparsec*
+
         """
         return self.transverse_comoving_distance(
             z2, params
@@ -189,13 +226,19 @@ def angular_diameter_distance(
         -----------
         z: Tensor
             The redshifts.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The angular diameter distance to each redshift.
+
+            *Unit: megaparsec*
+
         """
         return self.comoving_distance(z, params, **kwargs) / (1 + z)
 
@@ -210,15 +253,24 @@ def angular_diameter_distance_z1z2(
         ----------
         z1: Tensor
             The starting redshifts.
+
+            *Unit: unitless*
+
         z2: Tensor
             The ending redshifts.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The angular diameter distance between each pair of redshifts.
+
+            *Unit: megaparsec*
+
         """
         return self.comoving_distance_z1z2(z1, z2, params, **kwargs) / (1 + z2)
 
@@ -238,15 +290,24 @@ def time_delay_distance(
         ----------
         z_l: Tensor
             The lens redshifts.
+
+            *Unit: unitless*
+
         z_s: Tensor
             The source redshifts.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The time delay distance for each pair of lens and source redshifts.
+
+            *Unit: megaparsec*
+
         """
         d_l = self.angular_diameter_distance(z_l, params)
         d_s = self.angular_diameter_distance(z_s, params)
@@ -269,15 +330,24 @@ def critical_surface_density(
         ----------
         z_l: Tensor
             The lens redshifts.
+
+            *Unit: unitless*
+
         z_s: Tensor
             The source redshifts.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The critical surface density for each pair of lens and source redshifts.
+
+            *Unit: solMass/megaparsec^2*
+
         """
         d_l = self.angular_diameter_distance(z_l, params)
         d_s = self.angular_diameter_distance(z_s, params)
@@ -307,10 +377,19 @@ def __init__(
         Name of the cosmology.
         h0: Optional[Tensor]
             Hubble constant over 100. Default is h0_default.
+
+            *Unit: unitless*
+
         critical_density_0: (Optional[Tensor])
             Critical density at z=0. Default is critical_density_0_default.
+
+            *Unit: solMass/megaparsec^3*
+
         Om0: Optional[Tensor]
             Matter density parameter at z=0. Default is Om0_default.
+
+            *Unit: unitless*
+
         """
         super().__init__(name)
 
@@ -345,10 +424,15 @@ def hubble_distance(self, h0):
         h0: Tensor
             Hubble constant.
 
+            *Unit: unitless*
+
         Returns
         -------
         Tensor
             Hubble distance.
+
+            *Unit: megaparsec*
+
         """
         return c_Mpc_s / (100 * km_to_Mpc) / h0
 
@@ -370,13 +454,19 @@ def critical_density(
         ----------
         z: Tensor
             Redshift.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         torch.Tensor
             Critical density at redshift z.
+
+            *Unit: solMass/megaparsec^3*
+
         """
         Ode0 = 1 - Om0
         return critical_density_0 * (Om0 * (1 + z) ** 3 + Ode0)  # fmt: skip
@@ -393,10 +483,15 @@ def _comoving_distance_helper(
         x: Tensor
             Input tensor.
 
+            *Unit: megaparsec*
+
         Returns
         -------
         Tensor
             Computed comoving distances.
+
+            *Unit: megaparsec*
+
         """
         return interp1d(
             self._comoving_distance_helper_x_grid,
@@ -422,13 +517,19 @@ def comoving_distance(
         ----------
         z: Tensor
             Redshift.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             Comoving distance to redshift z.
+
+            *Unit: megaparsec*
+
         """
         Ode0 = 1 - Om0
         ratio = (Om0 / Ode0) ** (1 / 3)