docs: Update and add units to the docstrings (Ciela-Institute#162)

src/caustics/cosmology/base.py

@@ -17,12 +17,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: Mpc*
+
     Mass
-        solar mass
+
+        *Unit: Msun*
 
     Attributes
     ----------
@@ -50,13 +51,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: Msun/Mpc^3*
+
         """
         ...
 
@@ -72,13 +79,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: Mpc*
+
         """
         ...
 
@@ -94,13 +107,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: Mpc*
+
         """
         ...
 
@@ -115,15 +134,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: Mpc*
+
         """
         return self.comoving_distance(z2, params) - self.comoving_distance(z1, params)
 
@@ -138,15 +166,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: Mpc*
+
         """
         return self.transverse_comoving_distance(
             z2, params
@@ -163,13 +200,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: Mpc*
+
         """
         return self.comoving_distance(z, params, **kwargs) / (1 + z)
 
@@ -184,15 +227,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: Mpc*
+
         """
         return self.comoving_distance_z1z2(z1, z2, params, **kwargs) / (1 + z2)
 
@@ -212,15 +264,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: Mpc*
+
         """
         d_l = self.angular_diameter_distance(z_l, params)
         d_s = self.angular_diameter_distance(z_s, params)
@@ -243,15 +304,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: Msun/Mpc^2*
+
         """
         d_l = self.angular_diameter_distance(z_l, params)
         d_s = self.angular_diameter_distance(z_s, params)